#@input File(label="Working Directory", style="directory") expdir
#@input String(label="Channels", choices={"Red Channel Only", "Green Channel Only", "Red and Green Channels"}, style="listBox") channels
#@input String(label="Red Channel Prefix", value="RC_") rcpre
#@input String(label="Green Channel Prefix", value="GC_") gcpre
#@input String(label="Cell Mask Prefix", value="CM_") cmpre
#@input ImageJ ij

import os
import csv
from datetime import datetime
from net.imglib2.util import Intervals
from net.imglib2.type.numeric.integer import UnsignedByteType

RED_CHANNEL="Red Channel Only"
GREEN_CHANNEL="Green Channel Only"

imext = ".tif"

expdir = str(expdir)

with open(os.path.join(expdir, os.path.basename(expdir) + ".csv"), "w") as csvfile:

    output = csv.writer(csvfile)

    output.writerow(["Experiment", "Plate", "Well", "Field", "Timepoint", "Channel", "Value"])

    for root, dirs, files in os.walk(expdir):
        if not dirs:
            cms = []
            rcs = []
            gcs = []
            for f in files:
                if f.endswith(imext):
                    if f.startswith(cmpre):
                        cms.append(f[len(cmpre):-len(imext)])
                    if f.startswith(rcpre):
                        rcs.append(f[len(rcpre):-len(imext)])
                    if f.startswith(gcpre):
                        gcs.append(f[len(gcpre):-len(imext)])

            cms = sorted(set(cms))
            rcs = sorted(set(rcs))
            gcs = sorted(set(gcs))

            if channels == RED_CHANNEL:
                mismatch = cms != rcs
            elif channels == GREEN_CHANNEL:
                mismatch = cms != gcs
            else:
                mismatch = (cms != rcs) or (cms != gcs)

            if mismatch: raise Exception("No all necessary images found in %s" % root)

            bns = cms

            if len(bns) < 1:
                raise Exception("No images found in %s" % root)

            for bn in bns:
                info = bn.split("_")
                experiment = os.path.basename(expdir)
                plate = os.path.basename(root)
                timepoint = datetime.strptime(info[2]+info[3], "%Yy%mm%dd%Hh%Mm").isoformat()
                well = info[0]
                field = info[1]

                if channels == RED_CHANNEL:
                    chpres = [rcpre]
                elif channels == GREEN_CHANNEL:
                    chpres = [gcpre]
                else:
                    chpres = [rcpre, gcpre]

                for chpre in chpres:
                    if chpre == rcpre:
                        channel = "red"
                    elif chpre == gcpre:
                        channel = "green"

                    data = ij.io().open(os.path.join(expdir, root, chpre + bn + imext))
                    mask = ij.io().open(os.path.join(expdir, root, cmpre + bn + imext))

                    if not Intervals.equalDimensions(data, mask):
                        raise Exception("Dimensions of images do not match.")

                    mask = ij.op().run("threshold.apply", mask, UnsignedByteType(1))
                    mask = ij.op().run("convert.float32", mask)

                    data = ij.op().run("math.multiply", mask, data)
                    
                    value = ij.op().run("stats.sum", data).get() / ij.op().run("stats.sum", mask).get()

                    print("Experiment = %s, Plate = %s, Well = %s, Field = %s, Timepoint = %s, Channel = %s, Value = %s" % (experiment, plate, well, field, timepoint, channel, value))

                    output.writerow([experiment, plate, well, field, timepoint, channel, value])

