import numpy as np
import matplotlib.pyplot as plt
import re
import os
import pickle as pkl
from skimage.measure import block_reduce
import plastia.analysis.fit as plastia_fit
from AffineTransformationSolver import affine_fit
import brewer2mpl as cb

# Function library
def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0))
    return (cumsum[N:] - cumsum[:-N]) / N


def ensure_dir(f):
    d = os.path.dirname(f)
    if not os.path.exists(d):
        os.makedirs(d)


def plot_graded_line(ax, x, y, s=None):
    assert len(x)==len(y)
    n = len(x)
    if s is None:
        s = len(x)/10
    t=np.linspace(0, 1, np.size(x))**2
    # Segement plot and colour depending on T
    s = 5 # Segment length
    for i in range(0, n-s, s):
        ax.plot(x[i:i+s+1], y[i:i+s+1], color=(1-t[i]**2, 0.2, t[i]))

    range_x = np.max(x) - np.min(x)
    range_y = np.max(y) - np.min(y)
    view_range = 1.1*max((range_x, range_y))

    ax.set_xlim([np.min(x)-(view_range-range_x)/2.0, np.max(x)+(view_range-range_x)/2])
    ax.set_ylim([np.min(y)-(view_range-range_y)/2.0, np.max(y)+(view_range-range_y)/2])
    #ax.set_aspect('equal', 'datalim')


# Class library

class HdfDirHandler(object):

    def __init__(self, datevec, scan, **kwargs):
        self.datevec = datevec
        self.scan = scan
        self.base_directory = kwargs.pop('base_directory', "W:Runlog")
        self.HDF_base_dir = '%s\\%04d\\%02d\\%02d\\%04d%02d%02d-%04d\\' % (
                 self.base_directory, self.datevec[0], self.datevec[1], self.datevec[2],
                 self.datevec[0], self.datevec[1], self.datevec[2], self.scan)

        self.day_dir = '%s\\%04d\\%02d\\%02d' % (
                 self.base_directory, self.datevec[0], self.datevec[1], self.datevec[2])


class LatticeData(HdfDirHandler):
    channels = {'nw': 0, 'ne': 1, 'nufern': 2}
    remainders = {'nw': 1, 'ne': 0}

    def __init__(self, datevec, scan, lattice, **kwargs):
        """
        :param lattice: must be 'ne', 'nw', or 'nufern'
        The Important attributes generated by this constructor are self.camera_data and self.ss_moments
        """
        assert lattice in LatticeData.channels.keys()
        self.lattice = lattice

        # Lattice stability scans are done in NE - NW - NE ... order. For this reason, it is good to add measurements
        # to this measurement if they remainder == idx % divider
        self.remainder = kwargs.pop('remainder', LatticeData.remainders[self.lattice])
        self.divider = kwargs.pop('divider', 2)

        # Single site relevant settings
        self.threshold = kwargs.pop('threshold', 0.45)
        self.is_analyze_ss = kwargs.pop('is_analyze_ss', False)

        # Directory handling
        super(LatticeData, self).__init__(datevec, scan, **kwargs)

        print "Getting %s Data" % self.lattice

        # Initialize Camera Readings
        #regex = re.compile('out%04d%02d%02d \d{4}-%s' % (datevec[0], datevec[1], datevec[2],
        #                                                 LatticeData.channels[self.lattice]))
        regex = re.compile('log-%s.csv' % (LatticeData.channels[self.lattice]))
        files = os.listdir("%s\\beam_mon\\" % self.HDF_base_dir)

        cam_stability = []
        for file in files:
            if regex.match(file) is not None:
                cam_stability.append(file)
        assert len(cam_stability) < 2  # if more than 1 camera file found, not sure which to use!
        if len(cam_stability) == 0:
            print "No camera stability file found for channel %s" % self.lattice
        stability_file = "%sbeam_mon\\%s" % (self.HDF_base_dir, cam_stability[0]) if cam_stability[0] is not None else None
        self.camera_data = np.genfromtxt(stability_file, delimiter=',')

        # Initialize SS analyzed moments
        if os.path.exists("%s\\FitResults" % self.HDF_base_dir) and self.is_analyze_ss:
            result_regex = re.compile("fitresult_(\d{3})-\d{2}.pkl")
            atom_matrixes = []
            shot_nums = []
            for file in os.listdir("%s\\FitResults" % self.HDF_base_dir):
                if result_regex.match(file) is not None:
                    shot_num = int(result_regex.match(file).groups()[0])
                    ss_data = pkl.load(open("%s\\FitResults\\%s" % (self.HDF_base_dir, file), "rb"))
                    atom_matrix = ss_data['amps'] > self.threshold
                    atom_matrixes.append(atom_matrix)
                    shot_nums.append(shot_num)
            combined = zip(shot_nums, atom_matrixes)
            combined.sort(key=lambda combo: combo[0])
            sorted_shot_nums, sorted_atom_matrixes = zip(*combined)

            moments = []
            for idx, matrix in enumerate(sorted_atom_matrixes):
                reduced, fit = LatticeData.get_moments(matrix)
                if idx % self.divider == self.remainder:
                    moments.append((fit.initial_guess[3], fit.initial_guess[2]))
            self.ss_moments = zip(*moments)
        else:
            self.ss_moments = None

    def report(self):
        print "REPORT %s" % self.lattice

        std_dict = {}
        # Camera moments
        if self.camera_data is not None:
            std_dict['sigma_camera'] = (np.std(self.camera_data[:, 1]),
                                        np.std(self.camera_data[:, 2]))  # Tuple containing the std
            std_dict['mean_camera'] = (np.mean(self.camera_data[:, 1]),
                                       np.mean(self.camera_data[:, 2]))  # Tuple containing the mean

            print "Camera sigma = %s, %s px" % std_dict['sigma_camera']
            print "Camera mean = %s, %s px" % std_dict['mean_camera']
            std_dict['camera_data'] = self.camera_data

        # Single site
        if self.ss_moments is not None:
            std_dict['sigma_ss'] = (np.std(self.ss_moments[0]),
                                    np.std(self.ss_moments[1]))  # Tuple containing the std
            std_dict['mean_ss'] = (np.mean(self.ss_moments[0]),
                                   np.mean(self.ss_moments[1]))  # Tuple containing the mean
            std_dict['ss_moments'] = self.ss_moments

            print "Camera sigma = %s, %s px" % std_dict['sigma_ss']
            print "Camera mean = %s, %s px" % std_dict['mean_ss']

        pkl.dump(std_dict, open("%s\\stability_analysis_results_%s.pkl" %(self.HDF_base_dir, self.lattice), 'wb'))

    def plot_stability_data(self, ax_list, is_running_avg = False, running_mean_step = 10):
        """
        Plots data with a continuous gradient standing for time
        :param axis: 'ne' or 'nw'
        :param ax_list, list of 2 elements containing axes on which to draw
        :param is_running_avg: if you want to do a running avg,
        :param running_mean_step:  make this the running avg length (int)
        :return: 0
        """
        assert len(ax_list) == 2
        ax1 = ax_list[0]
        ax2 = ax_list[1]

        # Camera Data
        x_cam = self.camera_data[:, 1]
        y_cam = self.camera_data[:, 2]

        ax1.set_xlabel('pix (3 um)')
        if is_running_avg:
            x_cam = running_mean(x_cam, running_mean_step)
            y_cam = running_mean(y_cam, running_mean_step)
        plot_graded_line(ax1, x_cam, y_cam)
        ax1.set_title('%s_camera' % self.lattice)

        # Atom Data
        x_atom = self.ss_moments[0] if self.ss_moments is not None else np.zeros_like(x_cam)
        y_atom = self.ss_moments[1] if self.ss_moments is not None else np.zeros_like(y_cam)
        if is_running_avg:
            x_atom = running_mean(x_atom, running_mean_step)
            y_atom = running_mean(y_atom, running_mean_step)

        plot_graded_line(ax2, x_atom, y_atom)
        ax2.set_title('%s_ss_moments' % self.lattice)
        ax2.set_xlabel('site (0.56 um)')
        postscript = "run_avg" if is_running_avg else "no_run_avg"

    @staticmethod
    def get_moments(matrix, **kwargs):
        """
        :param matrix: matrix to be analyzed
        :param kwargs:
        :return: reduced matrix, fit for the matrix (which auto-computed the moments)
        """
        block_size = kwargs.pop('block_size', 1)
        reduced = block_reduce(matrix, block_size=(block_size, block_size), func=np.mean)
        fit = plastia_fit.FitFunctionGaussian2DRotation()
        fit.guess_matrix(reduced)
        if kwargs.pop('is_fit', False):
            fit.fit_matrix(reduced)
        return reduced, fit


class Measurement(HdfDirHandler):

    def __init__(self, datevec, scan, **kwargs):
        """
        Assumes that the scan directory contains the
        :param datevec: date of scan (year, month, day)
        :param scan: scan num on this day
        :param kwargs:
        """
        self.is_running_avg = kwargs.pop('is_running_avg', 10)
        self.threshold = kwargs.pop("threshold", 0.45)

        # Directory handling
        super(Measurement, self).__init__(datevec, scan, **kwargs)

        self.ne_lattice = LatticeData(self.datevec, self.scan, 'ne', threshold=self.threshold,
                                      base_directory=self.base_directory, divider=2, remainder=0)
        self.nw_lattice = LatticeData(self.datevec, self.scan, 'nw', threshold=self.threshold,
                                      base_directory=self.base_directory, divider=2, remainder=1)
        self.ne_lattice.report()
        self.nw_lattice.report()

    def pickle_this(self):
        pkl.dump(self, open("%s\\measurement.pkl" % self.HDF_base_dir, 'wb'))

    def plot(self):
        fig = plt.figure(figsize=(10, 10))
        ax_ne_cam = fig.add_subplot(221)
        ax_ne_ss = fig.add_subplot(222)
        ax_nw_cam = fig.add_subplot(223)
        ax_nw_ss = fig.add_subplot(224)
        self.ne_lattice.plot_stability_data([ax_ne_cam, ax_ne_ss])
        self.nw_lattice.plot_stability_data([ax_nw_cam, ax_nw_ss])
        postscript = ""
        plt.savefig("%s\\stability_%s.png" % (self.HDF_base_dir, postscript))

        fig = plt.figure(figsize=(10, 10))
        ax_ne_cam = fig.add_subplot(221)
        ax_ne_ss = fig.add_subplot(222)
        ax_nw_cam = fig.add_subplot(223)
        ax_nw_ss = fig.add_subplot(224)
        self.ne_lattice.plot_stability_data([ax_ne_cam, ax_ne_ss], is_running_avg=True, running_mean_step=10)
        self.nw_lattice.plot_stability_data([ax_nw_cam, ax_nw_ss], is_running_avg=True, running_mean_step=10)
        postscript = "running_avg"
        plt.savefig("%s\\stability_%s.png" % (self.HDF_base_dir, postscript))


class AffineAnalysis(object):

    def __init__(self, date, scan_list, **kwargs):
        self.ne_ss = []
        self.nw_ss = []
        self.ne_cam = []
        self.nw_cam = []
        self.base_directory = kwargs.pop('base_directory', "W:Runlog")

        for scan_num in scan_list:
            hdf = HdfDirHandler(date, scan_num, base_directory=self.base_directory)
            ne_dict = pkl.load(open("%s//stability_analysis_results_ne.pkl" % hdf.HDF_base_dir, 'rb'))
            nw_dict = pkl.load(open("%s//stability_analysis_results_nw.pkl" % hdf.HDF_base_dir, 'rb'))
            self.nw_cam.append(np.array(nw_dict['mean_camera']))
            self.ne_cam.append(np.array(ne_dict['mean_camera']))
            self.nw_ss.append(np.array(nw_dict['mean_ss']))
            self.ne_ss.append(np.array(ne_dict['mean_ss']))

        self.results_directory = "%s\\%s" % (HdfDirHandler(date, scan_list[0],
                                                           base_directory=self.base_directory).day_dir,
                                             "affine_analysis")
        ensure_dir(self.results_directory)

        # perform the fits
        self.transform_ne, self.transformed_ne = self.estimate_affine_parameters(self.ne_cam, self.ne_ss)
        self.plot_affine_fit(self.transformed_ne, self.ne_ss, is_save=True, name="ne_fit")

        self.transform_nw, self.transformed_nw = self.estimate_affine_parameters(self.nw_cam, self.nw_ss)
        self.plot_affine_fit(self.transformed_nw, self.nw_ss, is_save=True, name="nw_fit")


    @staticmethod
    def estimate_affine_parameters(cam, ss):
        """
        :param cam: camera data, list of tuples
        :param ss: single site data, list of tuples
        :return: transformation object, transformed points into lattice coordinates
        """
        assert len(cam) == len(ss)
        transform = affine_fit(cam, ss)
        transformed = []
        for s in cam:
            transformed.append(transform.Transform(s))
        print transform

        return transform, transformed

    def plot_affine_fit(self, transformed, ss, is_save = False, name = "no_name"):
        """
        :param transformed: transformed coordinates
        :param ss: single site analyzed coordinates
        :param is_save: whether to save the plot
        :param name: name to give it
        :return:
        """
        assert len(transformed) == len(ss)

        fig = plt.figure(figsize=(5, 5))
        ax1 = fig.add_subplot(111)

        colors = cb.get_map('Dark2', 'qualitative', len(ss)).mpl_colors
        multi_color_scatter_plot(ax1, transformed, name, "sites", marker='d', colors=colors)
        multi_color_scatter_plot(ax1, ss, name, "sites", colors=colors)
        if is_save:
            plt.savefig("%s\\%s.png" % (self.results_directory, name))
        else:
            plt.show()

    def plot_data(self):
        fig = plt.figure(figsize=(10, 11))
        ax1 = fig.add_subplot(221)
        ax2 = fig.add_subplot(222)
        ax3 = fig.add_subplot(223)
        ax4 = fig.add_subplot(224)

        multi_color_scatter_plot(ax1, self.ne_cam, "NE Cam", "pix (3 um)")
        multi_color_scatter_plot(ax2, self.ne_ss, "NE SS", "sites")
        multi_color_scatter_plot(ax3, self.nw_cam, "NW Cam", "pix (3 um)")
        multi_color_scatter_plot(ax4, self.nw_ss, "NW SS", "sites")
        plt.savefig('%s//affine_transform.png' % self.results_directory)
        plt.show()


def multi_color_scatter_plot(ax, data, title, x_label, marker = 'o', ms=15, colors=None):
    """
    :param ax:
    :param data: a list of tuples corresponding to points
    :return:
    """
    data_lists = zip(*data)
    x = data_lists[0]; y = data_lists[1]
    print "plotting %s datapoints" % len(data)
    for idx, point in enumerate(data):
        if colors == None:
            ax.plot(point[0], point[1], marker=marker, ms=ms)
        else:
            ax.plot(point[0], point[1], marker=marker, ms=ms, color = colors[idx])
    range_x = np.max(x) - np.min(x)
    range_y = np.max(y) - np.min(y)
    view_range = 1.1*max((range_x, range_y))
    ax.set_xlim([np.min(x)-(view_range-range_x)/2.0, np.max(x)+(view_range-range_x)/2])
    ax.set_ylim([np.min(y)-(view_range-range_y)/2.0, np.max(y)+(view_range-range_y)/2])
    ax.set_aspect('equal')
    ax.set_title(title)
    ax.set_xlabel(x_label)


def analyze_stability_scans():

    date = (2016, 9, 15)
    #scan_list = [52, 57, 59, 61, 62, 63, 64]
    scan_list = [52]

    for scan in scan_list:
        print "analyzing scan %s" % scan
        measurement = Measurement(date, scan, base_directory='W:RunLog')
        measurement.plot()


def plot_transformed_comparison(transformed, ss_moments, name, ax_list, is_running_avg = False, running_mean_step = 10):
        """
        Plots data with a continuous gradient standing for time

        :param axis: 'ne' or 'nw'
        :param ax_list, list of 2 elements containing axes on which to draw
        :param is_running_avg: if you want to do a running avg,
        :param running_mean_step:  make this the running avg length (int)
        :return: 0
        """
        assert len(ax_list) == 2
        ax1 = ax_list[0]
        ax2 = ax_list[1]

        # Camera Data
        x_cam = transformed[0]
        y_cam = transformed[1]

        ax1.set_xlabel('site (0.56 um)')
        if is_running_avg:
            x_cam = running_mean(x_cam, running_mean_step)
            y_cam = running_mean(y_cam, running_mean_step)

        plot_graded_line(ax1, x_cam, y_cam)
        ax1.set_title('%s_camera' % name)

        # Atom Data
        x_atom = ss_moments[0]
        y_atom = ss_moments[1]
        if is_running_avg:
            x_atom = running_mean(x_atom, running_mean_step)
            y_atom = running_mean(y_atom, running_mean_step)

        plot_graded_line(ax2, x_atom, y_atom)
        ax2.set_title('%s_ss_moments' % name)
        ax2.set_xlabel('site (0.56 um)')
        postscript = "run_avg" if is_running_avg else "no_run_avg"


def analyze_affine_transformation():
    date = (2016, 9, 15)
    #scans = [52, 57, 59, 61, 62, 63, 64]
    scans = [57, 59, 61, 62, 63, 64]
    affine = AffineAnalysis(date, scans)

    stability_scans = [52]
    #stability_scans = [57]
    for scan in stability_scans:
        print "analyzing scan %s" % scan
        hdf_dir = HdfDirHandler(date, scan, base_directory='Y:Runlog')
        ne_dict = pkl.load(open("%sstability_analysis_results_ne.pkl" % hdf_dir.HDF_base_dir, 'rb'))
        nw_dict = pkl.load(open("%sstability_analysis_results_nw.pkl" % hdf_dir.HDF_base_dir, 'rb'))
        transformed_ne, ss_moments_ne = apply_transform(affine.transform_ne, ne_dict)
        transformed_nw, ss_moments_nw = apply_transform(affine.transform_nw, nw_dict)

        fig = plt.figure(figsize=(10, 10))
        ax_ne_cam = fig.add_subplot(221)
        ax_ne_ss = fig.add_subplot(222)
        ax_nw_cam = fig.add_subplot(223)
        ax_nw_ss = fig.add_subplot(224)
        plot_transformed_comparison(transformed_ne, ss_moments_ne, 'NE', [ax_ne_cam, ax_ne_ss])
        plot_transformed_comparison(transformed_nw, ss_moments_nw, 'NW', [ax_nw_cam, ax_nw_ss])

        fig = plt.figure(figsize=(10, 10))
        ax_ne_cam = fig.add_subplot(221)
        ax_ne_ss = fig.add_subplot(222)
        ax_nw_cam = fig.add_subplot(223)
        ax_nw_ss = fig.add_subplot(224)
        print ss_moments_nw
        plot_transformed_comparison(transformed_ne, ss_moments_ne, 'NE', [ax_ne_cam, ax_ne_ss], is_running_avg=True)
        plot_transformed_comparison(transformed_nw, ss_moments_nw, 'NW', [ax_nw_cam, ax_nw_ss], is_running_avg=True)

        plt.show()


def apply_transform(transform, data_dict):
    camera_data = data_dict['camera_data']
    ss_moments = data_dict['ss_moments']
    camera_data = map(lambda x: (x[1], x[2]), camera_data)
    transformed = []
    for data_pt in camera_data:
        transformed.append(transform.Transform(data_pt))
        #transformed.append(data_pt)
    transformed = np.array(zip(*transformed))
    ss_moments = np.array(ss_moments)

    return transformed, ss_moments


if __name__ == "__main__":
    #analyze_stability_scans()
    analyze_affine_transformation()
    '''
    #analyze_affine_transformation()
    x = np.array([1, 2])
    s = np.array([[0, 1], [0, 1]])
    print np.dot(s, x)
    '''

