# -*- coding: utf-8 -*-
"""
Script for DLS Data processing.

@author: Alex
zuletzt modifiziert: 03.08.2018
"""
# import matplotlib
# matplotlib.use("Agg")

import re
import math
import fitfunctions as ff
from constants import KB
import utils as ut


try:
    import numpy as np
except ImportError:
    input("Numpy not installed.")
    exit

try:
    import matplotlib.pyplot as plt
except ImportError:
    input("Matplotlib not installed.")
    exit

try:
    import scipy.optimize as sco
except ImportError:
    input("Scipy not installed.")

try:
    import os
except ImportError:
    input("os module not installed.")

# Adjusts path separator depending on OS.
if os.name == "nt":
    PSEP = "\\"
elif os.name == "posix":
    PSEP = "/"
else:
    print("Operating System not recognized. Exiting.")
    exit()


def getConstVal(filename, Name):
    """
    Reads out constants from ALV autogenerated file.

    filename: Name of the file to read angle from.
    Name: The name of the constant the function should look for
    returns: constant value.
    """
    with open(filename, "r") as f:
        for line in f:
            if Name in line:
                lsplit = line.split()
                for el in lsplit:
                    if ut.isFloat(el):
                        return float(el)
    print("Nothing found.")
    return


def getSampleName(filename):
    """
    Reads out samplename from ALV autogenerated file.

    filename: Name of the file to read angle from.
    """
    with open(filename, "r") as f:
        for line in f:
            if "Samplename" in line:
                lsplit = line.split()
                return lsplit[2][1:-1]


def getDynData(filename):
    """
    Reads out delay time, Autocorrelation function and standard deviation

    returns: list of delay times tau, AKF, Measurement Time and Countrate
    """
    flag1 = False
    flag2 = False
    delayL = []
    akfL = []
    measTimeL = []
    countrateL = []

    with open(filename, "r") as f:
        for line in f:
            # findet dynamische Werte, da sie nach "Correlation" auftreten.
            if flag1:
                corrVals = line.split()
                try:
                    if ut.isFloat(corrVals[0]) and ut.isFloat(corrVals[1]):
                        delayL.append(float(corrVals[0]))
                        akfL.append(float(corrVals[1]))
                except IndexError:
                    pass
            if "Correlation" in line:
                flag1 = True

            if flag2:
                timeAndCountrate = line.split()
                try:
                    if ut.isFloat(timeAndCountrate[0]):
                        measTimeL.append(float(timeAndCountrate[0]))
                except IndexError:
                    pass

                try:
                    if ut.isFloat(timeAndCountrate[1]):
                        countrateL.append(float(timeAndCountrate[1]))
                except IndexError:
                    pass
            if "Count Rate" in line:
                flag1 = False
                flag2 = True

    return delayL, akfL, measTimeL, countrateL


class ALVData(object):
    """
    The corresponding data of a single ALV generated file.
    Contains various information gathered by getDynData and getConstVal.

    filename: Name of the file the data is located.
    samplename = Name of the sample.
    angle = Angle in degrees under which the sample was measured.
    refInd = refractive index of the sample.
    delay = Delay time tau of the autocorrelation function (akf)
    akf = Autocorrelation function calculated by ALV software (g2 - 1)
    stdDev = Standard deviation of the autocorrelation function.
    """

    def __init__(self, filename, fitfun):
        dyndata = getDynData(filename)

        self.filename = filename
        self.fitfun = fitfun
        self.samplename = getSampleName(filename)
        self.angle = getConstVal(filename, "Angle")
        self.refInd = getConstVal(filename, "Refractive Index")
        self.wavelength = getConstVal(filename, "Wavelength") * 1e-9
        self.visc = getConstVal(filename, "Viscosity")
        self.temp = getConstVal(filename, "Temperature")
        self.delay = np.array(dyndata[0])
        self.akf = np.array(dyndata[1])
        self.meastime = np.array(dyndata[2])
        self.countrate = np.array(dyndata[3])
        self.coeffs = None
        self.A = None
        self.q = 4*np.pi * self.refInd / self.wavelength * \
            np.sin(np.deg2rad(self.angle/2))
        self.meanCR = getConstVal(filename, "MeanCR0")

    def __str__(self):
        return self.filename

    def __repr__(self):
        return "ALVData("+self.filename+")"

    def curveFit(self, delay, akf):
        """
        Uses fitFun to plot data to single exponential autocorrelation function
        """
        popt, pcov = sco.curve_fit(self.fitfun, delay, akf)
        self.coeffs = popt[:]
        return popt, pcov

    def saveAKF(self, filepath, start=0, fallOff=0):
        """
        Saves Plot of the autocorrelation function versus the delay time tau
        in a file, fitted in a single exponential. A dynamic datarange can
        be specified using the start and fallOff parameters.

        filepath: Path of the file where the plot should be saved to.
        datarange: Range of how many data points in the array should be used.
        start: Only uses point in the array after this value.
        fallOff: Slices array after start-value/current-value is lower than the
        fallOff.
        """
        # schneidet Delay und Autokorrelation nach den Vorgaben der datarange
        startValue = self.akf[start]
        akf = []

        for el in self.akf[start:]:
            if el/startValue > fallOff:
                akf.append(el)
            else:
                break

        delay = self.delay[start:len(akf)+start]

        plt.figure(self.filename, dpi=100)
        plt.clf()
        plt.semilogx(self.delay, self.akf, " bo", markersize=2,
                     label="Autocorrelation data")
        popt, pcov = self.curveFit(delay, akf)
        plt.semilogx(delay, self.fitfun(delay, *popt), '-r',
                     label="%s fit" % ff.nameOf(self.fitfun))
        plt.legend()
        plt.xlabel(r"Delay Time " + chr(964) + " [ms]")
        plt.ylabel("Autocorrelation Function g2 - 1")
        plt.title(self.samplename + ", Angle " + str(self.angle) + "°")

        plt.savefig(filepath)
        plt.close()
        return popt, pcov

    def createDir(self):
        """
        Creates a directory for the current unique sample, if the directory
        doesnt already exist.

        returns: path of the directory
        """
        # Strips the samplename with .txt ending off the filename, attaches the
        # samplename back again.
        targetPath = re.sub(self.samplename+r"(\d+)\.(\w)+$",
                            self.samplename, self.filename)

        if not os.path.exists(targetPath):
            os.makedirs(targetPath)
            return targetPath
        else:
            return targetPath


class ALVSample(object):
    """
    Represent one sample, consisting of a number of measurements.
    These measurements must consist of ALVData objects. The ALV Sample can then
    be represented via different plotting and logging mechanisms to the outside
    world.

    refInd: refractive Index
    wavelength: Wavelength
    visc: Viscosity
    temp: Temperature
    samplename: samplename (every sample added to the sample must be the same!)
    savingPath: path where to save the plots
    """

    def __init__(self, refInd, wavelength, visc, temp, samplename, savingPath):
        self.refInd = refInd
        self.wavelength = wavelength
        self.visc = visc
        self.temp = temp
        self.samplename = samplename
        self.path = savingPath
        self.dataObjs = []
        self.plotmode = plt.plot
        self.angles = []
        self.meanCRs = []
        self.qSqu = []
        self.coeffs = []

    def addData(self, data):
        """
        Adds ALV Data objects to the samples. Accepts lists or single objects.
        """
        if type(data) is not list:
            data = [data]
        for el in data:
            if not isinstance(el, ALVData):
                raise("Error in ALVSample: Data given is not an ALV Object.")
            self.dataObjs.append(el)
            self.angles.append(el.angle)
            self.meanCRs.append(el.meanCR)
            self.qSqu.append(el.q**2)
            self.coeffs.append(el.coeffs)
        return self.dataObjs

    def changePlotMode(self, plotmode):
        """
        Changes plotting mode, for example to semilogx or semilogy.
        """
        self.plotmode = plotmode
        return self.plotmode

    def plotMeanCRsDLS(self):
        """
        Plots mean countrate over angle, extracted from the given filenames.
        """

        plt.figure(self.samplename, dpi=100)
        plt.clf()
        self.plotmode(self.angles, self.meanCRs,
                      ' bo', markersize=2, label="meanCR")
        plt.legend()
        plt.title(self.samplename + ", meanCR over angle")
        plt.xlabel(r"$\theta$ in °")
        plt.ylabel("Mean Countrate")
        plt.savefig(self.path + PSEP
                    + self.samplename + "meanCR.png")
        plt.close()
        return

    def plotCoherenceFactor(self):
        """
        Plots Coherence Factor (first coefficient in fit function) over q^2.

        Works for fitfunctions cum2, cum3, singleExp.
        """
        # fit function not relevant for countrate, just take the simplest one

        coherenceFactor = [coeff[0] for coeff in self.coeffs]

        plt.figure(self.samplename, dpi=100)
        plt.clf()
        self.plotmode(self.angles, coherenceFactor, ' bo', markersize=2,
                      label="meanCR")
        plt.legend()
        plt.title(self.samplename + ", Coherence Factor over angle")
        plt.xlabel(r"$\theta$ in °")
        plt.ylabel("Coherence Factor")
        plt.savefig(self.path + PSEP
                    + self.samplename + "CF.png")
        plt.close()
        return

    def plotGammaAndFit(self):
        """
        Plots gamma over q^2 (first coefficient in fit function) over q^2.

        Works for fitfunctions cum2, cum3, singleExp.
        """
        gamma = [coeff[1] for coeff in self.coeffs]
        qSqu = np.array(self.qSqu)

        def linfun(x, a): return a*x

        popt, pcov = sco.curve_fit(linfun, qSqu, gamma)

        plt.figure(self.samplename, dpi=100)
        plt.clf()
        self.plotmode(qSqu, gamma, ' bo', markersize=2, label="meanCR")

        self.plotmode(qSqu, linfun(qSqu, *popt), 'r', label="Linear fit")
        plt.legend()
        plt.title(self.samplename + ", Gamma over q^2")
        plt.xlabel(r"q^2 in 1/m^2")
        plt.ylabel("Gamma")
        plt.savefig(self.path + PSEP
                    + self.samplename + "Gamma.png")
        plt.close()

        hydrodynR = ut.getHydroDynR(popt[0], self.visc, self.temp)
        return hydrodynR

    def plotHydrodynRadius(self):
        """
        Plots Hydrodynamic radius over countrate.

        D_s = Gamma/q**2
        """
        visc = self.visc
        temp = self.temp
        gamma = [coeff[1] for coeff in self.coeffs]
        qSqu = np.array(self.qSqu)
        print(self)
        hydroR = [1e9*ut.getHydroDynR(gamma[i]/qSqu[i], visc, temp)
                  for i in range(len(qSqu))]

        plt.figure(self.samplename, dpi=100)
        plt.clf()
        self.plotmode(self.angles, hydroR, ' bo', markersize=2,
                      label="Hydrodynamic radius")

        plt.legend()
        plt.title(self.samplename + ", Hydrodynamic radius over angle")
        plt.xlabel(r"$\theta$ in °")
        plt.ylabel("Hydrodynamic Radius in nm")
        plt.savefig(self.path + PSEP
                    + self.samplename + "hydroR.png")
        plt.close()

        return


def compareCRs(filenames1, filenames2):
    """
    Compares two mean countrates (diving 1 by 2 (1/2)) and plots them.
    """
    sample1 = readElementsAndProcess(filenames1, ff.singleExp,
                                     saveImg=False)
    sample2 = readElementsAndProcess(filenames2, ff.singleExp,
                                     saveImg=False)
    CRs1 = sample1.meanCRs
    CRs2 = sample2.meanCRs
    divCR = []
    for i in range(len(CRs1)):
        divCR.append(CRs1[i]/CRs2[i])

    plt.figure(sample1.samplename, dpi=100)
    plt.clf()
    sample1.plotmode(sample1.angles, divCR, ' bo', markersize=2,
                     label="Mean CR divided")

    plt.legend()
    plt.title(sample1.samplename + " divided by "
              + sample2.samplename)
    plt.xlabel(r"Angle in °")
    plt.ylabel("Countrate")
    plt.savefig(sample1.path + PSEP
                + sample1.samplename + "By" + sample2.samplename
                + "dividedCR.png")
    plt.close()
    return


def readElementsAndProcess(filenames, fitfun, start=0, fallOff=0,
                           saveImg=True):
    """
    Simply combines basic execution of ALV Elements functions. Can take any
    amount of filenames in an array.

    The possibility in saveAKF to give out a data range can be leveraged using
    the start and fallOff parameters as indicated in the saveAKF function.

    filenames: list of filenames that correspond to ALV Data, alternatively a
    single string that will be converted into a list. If the path to a direc-
    tory is given, the files inside the directory will be used as a file list.
    fitfun: fitfunction, can be supplied using the fitfunction.py
    start: array value at which fitting begins.
    fallOff: quotient at which fitting stops.
    log: Set to true if coefficients, countrate, angle and further stats should
         be logged.
    saveImg: Set to true if the fitfunctions should be plotted together with
             the data in a separate folder.

    returns dictionary of data about the whole sample for further processing.
    """
    # enables the function to also accept a single string as input, which
    # can be a directory.

    if type(filenames) is str:
        if os.path.isdir(filenames):
            contents = ut.getFilesInFolder(filenames, r"\d+\.\w+$")
            filenames = [filenames + PSEP
                         + name for name in contents]
            print(filenames)
        else:
            filenames = [filenames]

    if len(filenames) == 0:
        print("Empty List passed into function")
        return

    # reciproke scattering vector
    alvDataElements = []

    # indicates the return statement that at least one valid file was found
    # (or none).
    oneFileFoundFlag = False

    for data in filenames:

        # prepares data to be returned in dictionary later
        try:
            el = ALVData(data, fitfun)
        except:
            print("Could not open " + data + ".")
            continue

        alvDataElements.append(el)

        startValue = el.akf[start]
        akf = []

        for k in el.akf[start:]:
            if k/startValue > fallOff:
                akf.append(k)
            else:
                break
        delay = el.delay[start:len(akf)+start]

        oneFileFoundFlag = True
        el.curveFit(delay, akf)
        savingPath = el.createDir()

        if saveImg:
            impath = savingPath + PSEP + el.samplename + str(int(el.angle)) \
                + "Grad.png"
            el.saveAKF(impath, start, fallOff)

    if not oneFileFoundFlag:
        print("No valid file was found by the algorithm.")
        return None

    # initializes ALV Sample to be returned
    first = alvDataElements[0]
    sample = ALVSample(first. refInd, first.wavelength, first.visc, first.temp,
                       first.samplename, savingPath)
    sample.addData(alvDataElements)
    return sample


def dlsplot(filenames, fitfun, start=0, fallOff=0, log=False,
            noPlot=False, plotCorr=True, plotMeanCR=True,
            plotCoherence=True, plotGamma=True, plotHydroR=True,
            plotmode=plt.plot):
    """
    Wraps up all plotting functions into one interface.
    Attention: Setting no-plot to true overrides all plotting options.
    """

    if noPlot:
        plotCorr = False
        plotMeanCR = False
        plotCoherence = False
        plotGamma = False

    sample = readElementsAndProcess(filenames, fitfun,
                                    start, fallOff, plotCorr)
    sample.changePlotMode(plotmode)

    print(sample.coeffs)
    if plotMeanCR:
        sample.plotMeanCRsDLS()
    if plotCoherence:
        sample.plotCoherenceFactor()
    if plotGamma:
        rad = sample.plotGammaAndFit()
        print("Hydrodynamic Radius by linear fit:", rad)
    if plotHydroR:
        sample.plotHydrodynRadius()

    # Logging Mechanism: Prepares results for reading out via origin/other
    # data evaulation programs
    if log:
        txpath = sample.path + PSEP + sample.samplename + "log.asc"
        with open(txpath, "w+") as txfile:

            # Writes Header, containing information about sample
            angles = ut.list2string(sample.angles)
            qSqu = ut.list2string(sample.qSqu)

            columnTitles = "Delay Time\t" + "Autocorrelation g2-1\n"
            txfile.write(columnTitles)

            comments = "# Parameters: Wavelength Viscosity Temperature" \
                       + " Ref. Ind Angles q^2\n"

            txfile.write(comments)

            header = str(sample.wavelength) + "\n" \
                                            + str(sample.visc) + "\n" \
                                            + str(sample.temp) + "\n" \
                                            + str(sample.refInd) + "\n"
            header += "\\\t" + "\t".join(angles) + "\n" \
                      + "\\\t" + "\t".join(qSqu) + "\n"
#                     + "q^2\t" + ut.list2string(sample.qSqu) + "\n"

            txfile.write(header)

            # assumes delay times are equal for every data entry in the sample
            delayTimes = ut.list2string(sample.dataObjs[0].delay)
            for i in range(len(delayTimes)):
                entry = [str(data.akf[i]) for data in sample.dataObjs]
                txfile.write(delayTimes[i] + "\t" + "\t".join(entry) + "\n")

        # writes mean Countrate data to separate file
        pathmeanCR = sample.path + PSEP + sample.samplename + "CR.asc"
        with open(pathmeanCR, "w+") as crfile:
            angles = ut.list2string(sample.angles)
            columnTitles = "\t".join(angles)+"\n"
            crfile.write(columnTitles)
            print(sample.meanCRs)
            crs = ut.list2string(sample.meanCRs)
            crfile.write("\t".join(crs) + "\n")

    return sample
