import pandas as pd
from regmaxsn.core.occupancyBasedMeasure import occupancyEMD
import pathlib2
import json
import numpy as np
from matplotlib import pyplot as plt
import seaborn as sns
from regmaxsn.core.matplotlibRCParams import mplPars
import sys


dirPath = pathlib2.Path("/media/ajay/ADATA_HD720/Ginjang/DataAndResults/morphology/OriginalData/Tests")
expName = 'HSN-fluoro01.CNG'
swcSetSize = 20
N = 500

resDir = pathlib2.Path("/media/ajay/ADATA_HD720/Ginjang/DataAndResults/morphology/OccupancyBaseTests")
if not resDir.exists():
    resDir.mkdir()
voxelSize = 10

scaleDF = pd.DataFrame()
rotDF = pd.DataFrame()
transDF = pd.DataFrame()

suffixes = ("RandRotY", "RandScaleY", "RandTranslateY")
dfs = [rotDF, scaleDF, transDF]
labels = ("Rotation (degrees)", "Scale", "Translation (um)")
jsonKeys = ("angles", "scale", "translation")
parInds = (1, 1, 1)
figs = []


def saveData():
    for suffixInd, suffix in enumerate(suffixes):

        print("Doing {}".format(labels[suffixInd]))

        parSWCDict = {}

        for ind in range(N):
            label = labels[suffixInd]
            jsonKey = jsonKeys[suffixInd]
            parInd = parInds[suffixInd]
            outFile = str(dirPath / '{}{}{}.swc'.format(expName, suffix, ind))
            transJSONFile = str(dirPath / '{}{}{}.json'.format(expName, suffix, ind))
            with open(transJSONFile, "r") as fle:
                transJSON = json.load(fle)
            jsonPars = transJSON[jsonKey]
            parSWCDict[jsonPars[1]] = outFile

        allPars = parSWCDict.keys()
        allParsSorted = np.sort(allPars)

        maxStepSize = int(np.floor(float(N) / float(swcSetSize)))
        baseSet = np.arange(0, swcSetSize)
        for stepSize in range(1, maxStepSize + 1):
            print("Doing StepSize {}/{}".format(stepSize, maxStepSize))
            windowSlideSize = int(stepSize * swcSetSize / 2)
            windowStarts = range(0, N - stepSize * swcSetSize + 1, windowSlideSize)
            for windowStart in windowStarts:
                print("Doing Window start {}/{}".format(windowStart, windowStarts))
                pars = allParsSorted[windowStart + stepSize * baseSet]
                swcFiles = [parSWCDict[par] for par in pars]
                metric = occupancyEMD(swcFiles, voxelSize)
                if suffix == "RandRotY":
                    pars = np.rad2deg(pars)
                tempDict = {"mean of {}".format(label): np.mean(pars),
                            "std of {}".format(label): np.std(pars),
                            "metric": metric}
                dfs[suffixInd] = dfs[suffixInd].append(tempDict, ignore_index=True)

        outFile = str(resDir / "metricVs{}.xlsx".format(suffix))
        dfs[suffixInd].to_excel(outFile)

def plotData():

    sns.set(rc=mplPars)
    figs = []
    for suffixInd, suffix in enumerate(suffixes):
        label = labels[suffixInd]
        dfXL = str(resDir / "metricVs{}.xlsx".format(suffix))
        df = pd.read_excel(dfXL)
        df["mean of {}".format(label)] = df["mean of {}".format(label)].apply(lambda x: round(x, 3))
        df["std of {}".format(label)] = df["std of {}".format(label)].apply(lambda x: round(x, 3))
        df2Plot = df.pivot(index="mean of {}".format(label),
                           columns="std of {}".format(label),
                           values="metric")
        fig, ax = plt.subplots(figsize=(14, 11.2))
        sns.heatmap(data=df2Plot, ax=ax, xticklabels=10, yticklabels=10, cmap=plt.cm.jet)
        ax.set_xticklabels(ax.get_xticklabels(), rotation="vertical")
        ax.set_yticklabels(ax.get_yticklabels(), rotation="horizontal")
        fig.tight_layout()
        outFile = str(resDir / "metricVs{}.png".format(suffix))
        fig.savefig(outFile, dpi=150)
        figs.append(fig)
    plt.show()


if __name__ == "__main__":

    assert len(sys.argv) == 2, 'Improper usage! Please use as \'python {arg} save\' or' \
                               '\'python {arg} plot'.format(arg=sys.argv[0])

    if sys.argv[1] == "save":
        saveData()
    elif sys.argv[1] == "plot":
        plotData()
    else:
        raise ValueError