#!/usr/bin/python3

import sys
import pyvips
import numpy as np
from numpy import linalg

stain = [
    [0.468, 0.023, 0.767],
    [0.721, 0.141, 0.576],
    [0.511, 0.990, 0.284]
]
stain_inverse = linalg.inv(stain).tolist()

image = pyvips.Image.new_from_file(sys.argv[1])

# transform to linear 0-1
image = image.colourspace("scrgb")

# to stain space
stain_space = image.log().recomb(stain_inverse)
n_stains = stain_space.bands

# split by stain
stains = [stain_space.recomb(pyvips.Image.new_from_array(
            [[0] * n_stains] * i + 
            [[0] * i + [1] + [0] * (n_stains - i - 1)] + 
            [[0] * n_stains] * (n_stains - i - 1)))
          for i in range(n_stains)]

# from stain space back to a set of separated scrgb Images
images = [x.recomb(stain).exp() for x in stains]

# back to srgb
images = [x.colourspace("srgb") for x in images]

for i in range(len(images)):
    filename = f"stain-{i}.png"
    print(f"writing {filename} ...")
    images[i].write_to_file(filename)