from __future__ import annotations
from scikit_fluorescence import typing

import numpy as np
import json
import yaml
import pathlib

import scikit_fluorescence as skf
import scikit_fluorescence.io.ascii
#import scikit_fluorescence.decay


# Try to import pylab for plotting of data
try:
    import pylab as plt
except ImportError:
    plt = None

#from scikit_fluorescence.decay.tcspc import FluorescenceDecay


def read_tcspc_csv(
        filename: str = None,
        skiprows: int = None,
        rebin: typing.Tuple[int, int] = (1, 1),
        dt: float = 1.0,
        matrix_columns: typing.Tuple[int, int] = (0, 1),
        use_header: bool = False,
        is_jordi: bool = False,
        polarization: str = "vm",
        g_factor: float = 1.0,
        infer_delimiter: bool = True,
        *args,
        **kwargs
) -> typing.List[skf.decay.tcspc.FluorescenceDecay]:
    """

    :param filename:
    :param skiprows:
    :param rebin:
    :param dt:
    :param matrix_columns:
    :param use_header:
    :param is_jordi:
    :param polarization:
    :param g_factor:
    :param setup:
    :param args:
    :param kwargs:
    :return:
    """
    verbose = kwargs.get('verbose', skf.VERBOSE)
    if verbose:
        print('READ_TCSPC_CSV')
        print('-- filename:', filename)
        print('-- skiprows:', skiprows)
        print('-- rebin:', rebin)
        print('-- dt:', dt)
        print('-- matrix_columns:', matrix_columns)
        print('-- use_header:', use_header)
        print('-- is_jordi:', is_jordi)
        print('-- polarization:', polarization)
        print('-- g_factor:', g_factor)
        print('-- infer_delimiter:', infer_delimiter)
    rebin_x, rebin_y = rebin
    if is_jordi:
        infer_delimiter = False
        mc = None
    else:
        mc = matrix_columns
    csv_loader = skf.io.ascii.Csv(*args, **kwargs)
    csv_loader.load(
        filename,
        skiprows=skiprows,
        use_header=use_header,
        usecols=mc,
        infer_delimiter=infer_delimiter
    )
    data = csv_loader.data
    if is_jordi:
        if data.ndim == 1:
            data = data.reshape(1, len(data))

        n_data_sets, n_vv_vh = data.shape
        n_data_points = n_vv_vh // 2
        c1, c2 = data[:, :n_data_points], data[:, n_data_points:]

        new_channels = int(n_data_points / rebin_y)
        c1 = c1.reshape([n_data_sets, new_channels, rebin_y]).sum(axis=2)
        c2 = c2.reshape([n_data_sets, new_channels, rebin_y]).sum(axis=2)
        n_data_points = c1.shape[1]

        if polarization == 'vv':
            y = c1
            ey = skf.decay.tcspc.counting_noise(c1)
        elif polarization == 'vh':
            y = c2
            ey = skf.decay.tcspc.counting_noise(c2)
        elif polarization == 'vv/vh':
            e1 = skf.decay.tcspc.counting_noise(c1)
            e2 = skf.decay.tcspc.counting_noise(c2)
            y = np.vstack([c1, c2])
            ey = np.vstack([e1, e2])
        else:
            f2 = 2.0 * g_factor
            y = c1 + f2 * c2
            ey = skf.decay.tcspc.combined_counting_noise_parallel_perpendicular(
                parallel=c1,
                perpendicular=c2,
                g_factor=g_factor
            )
        x = np.arange(n_data_points, dtype=np.float64) * dt
    else:
        try:
            # check if time axis is channel number
            x = data[0]
            y = data[1:]
            if x[1].is_integer():
                info = "Axis in file seems to be a channel numbers. Multiplying axis by %s" % dt
                skf.logging.info(info)
                if verbose:
                    print(info)
                x *= dt
        except IndexError:
            info = "Could not reading file: {}".format(filename)
            skf.logging.warning(info)
            print(info)
            for line in open(filename, 'r').readlines()[:50]:
                print(line)
        n_datasets, n_data_points = y.shape
        n_data_points = int(n_data_points / rebin_y)
        try:
            y = y.reshape(
                [n_datasets, n_data_points, rebin_y]
            ).sum(axis=2)
            ey = skf.decay.tcspc.counting_noise(y)
            x = np.average(
                x.reshape([n_data_points, rebin_y]), axis=1
            ) / rebin_y
        except ValueError:
            print("Cannot reshape array")

    x = x[n_data_points % rebin_y:]
    y = y[n_data_points % rebin_y:]

    # rebin along x-axis
    y_rebin = np.zeros_like(y)
    ib = 0
    for ix in range(0, y.shape[0], rebin_x):
        y_rebin[ib] += y[ix:ix+rebin_x, :].sum(axis=0)
        ib += 1
    y_rebin = y_rebin[:ib, :]
    data_curves = list()
    for i, yi in enumerate(y_rebin):
        eyi = ey[i]
        # check for nans in beginning of decay
        offset = 0
        for i, n in enumerate(x):
            if np.isnan(n):
                offset = i + 1

        data: skf.decay.tcspc.FluorescenceDecay = {
            'time_axis': x[offset:],
            'decay_histogram': yi[offset:],
            'counting_noise': eyi[offset:]
        }
        data_curves.append(data)
    return data_curves


def process_tcspc_yaml_measurement_id(
        experiment_file: pathlib.Path,
        measurement_id: str,
        experiment_dict: dict,
        data_idx: int = 0,
        verbose: bool = skf.VERBOSE
) -> skf.decay.tcspc.TCSPCDataset:
    def read_decay(
            experiment_file: pathlib.Path,
            measurement_id: str,
            experiment_dict: dict,
            data_idx: int = 0,
            verbose: bool = skf.VERBOSE
    ) -> skf.decay.tcspc.FluorescenceDecay:
        measurement = experiment_dict['Measurement datasets'][measurement_id]
        filename = experiment_file.parent / measurement['Data']['Filename']
        setup_setting_name = measurement['Setup settings']
        setup_name = experiment_dict['Setup settings'][setup_setting_name]['Setup name']
        setup = experiment_dict['Setup'][setup_name]
        setup_part_ids = setup['Setup part IDs']
        setup_parts = experiment_dict['Setup parts']
        dt = 1.0
        for setup_part_id in setup_part_ids:
            setup_part = setup_parts[setup_part_id]
            if 'Counting electronics channel width' in setup_part.keys():
                dt = setup_part['Counting electronics channel width']
                break
        if verbose:
            skf.logging.info("READ_DECAY")
            skf.logging.info("-- measurement: %s" % measurement)
            skf.logging.info("-- filename: %s" % filename)
            skf.logging.info("-- setup_setting_name: %s" % setup_setting_name)
            skf.logging.info("-- setup_name: %s" % setup_name)
            skf.logging.info("-- dt: %s" % dt)
        data = skf.io.tcspc.read_tcspc_csv(
            filename,
            dt=dt,
            verbose=verbose
        )
        return data[data_idx]

    if verbose:
        print('PROCESS_MEASUREMENT_ID')
        print('-- experiment_file', experiment_file)
        print('-- measurement_id', measurement_id)
        print('-- data_idx', data_idx)
    measurement = experiment_dict['Measurement datasets'][measurement_id]
    data = read_decay(
        experiment_file=experiment_file,
        experiment_dict=experiment_dict,
        measurement_id=measurement_id,
        data_idx=data_idx
    )
    try:
        irf = read_decay(
            experiment_file=experiment_file,
            experiment_dict=experiment_dict,
            measurement_id=measurement['Instrument response function']['Measurement ID'],
            data_idx=data_idx
        )
        irf_hist = irf['decay_histogram']
    except KeyError:
        irf_hist = np.zeros(data['decay_histogram'])
        irf_hist[0] = 1.0
    try:
        lin_key = measurement['Linearization']['Measurement ID']
    except KeyError:
        info = '-- No linearization for measurement_id: %s' % measurement_id
        skf.logging.info(info)
        if verbose:
            print(info)
        lin_key = None
    if lin_key:
        lin = read_decay(
            experiment_file=experiment_file,
            experiment_dict=experiment_dict,
            measurement_id=lin_key,
            data_idx=data_idx,
            verbose=verbose
        )
        lin_hist_raw = lin['decay_histogram']
        try:
            lin_settings = measurement['Linearization']['Settings']
            lin_hist = skf.decay.tcspc.compute_linearization_table(
                data=lin_hist_raw,
                **lin_settings
            )
        except KeyError:
            info = "No linearization settings provided. Using raw data."
            skf.logging.warning(info)
            if verbose:
                print(info)
            lin_hist = lin_hist_raw
    else:
        lin_hist = np.ones_like(data['decay_histogram'])
    decay: skf.decay.tcspc.TCSPCDataset = {
        'time_axis': data['time_axis'],
        'decay_histogram': data['decay_histogram'],
        'counting_noise': data['counting_noise'],
        'instrument_response_function': irf_hist,
        'linearization': lin_hist
    }
    return decay


def read_tcspc_yaml(
        experiment_file: pathlib.Path,
        flat_dict: bool = False,
        verbose: bool = skf.VERBOSE
) -> typing.Dict:
    if verbose:
        skf.logging.info('READ_EXPERIMENT_YAML')
        skf.logging.info('-- experiment_file', experiment_file)
    with open(experiment_file, 'r') as fp:
        experiment_dict = yaml.safe_load(fp)
        measurement_id = experiment_dict['Sample']['Measurement ID']
        decay = process_tcspc_yaml_measurement_id(
            measurement_id=measurement_id,
            experiment_file=experiment_file,
            experiment_dict=experiment_dict,
            verbose=verbose
        )
        re = {
            'decay': decay,
            'references': {
            }
        }
        if not flat_dict:
            for reference_name in experiment_dict['Sample']['Reference']:
                reference_id = experiment_dict['Sample']['Reference'][reference_name]
                re['references'][reference_name] = process_tcspc_yaml_measurement_id(
                    measurement_id=reference_id,
                    experiment_dict=experiment_dict,
                    experiment_file=experiment_file,
                    verbose=verbose
                )
        else:
            re = {
                'decay': decay,
            }
            for reference_name in experiment_dict['Sample']['Reference']:
                reference_id = experiment_dict['Sample']['Reference'][reference_name]
                re[reference_name] = process_tcspc_yaml_measurement_id(
                    measurement_id=reference_id,
                    experiment_dict=experiment_dict,
                    experiment_file=experiment_file,
                    verbose=verbose
                )

    return re


def read_dataset_from_fps_json(
        json_file: str,
        verbose: bool = skf.VERBOSE,
        plot_data: bool = False
) -> typing.Dict[str, typing.Dict]:
    """

    Parameters
    ----------
    json_file
    verbose
    plot_data

    Returns
    -------
    Dict:
    experiment_name (key in JSON file): experiment data (dict)

    """
    def plot_data_sets(
            data_sets: typing.Dict
    ):
        if p is None:
            skf.logging.warning("Cannot plot. matplotlib not installed?")
            return None
        # plot the data
        for data_key in list(data_sets.keys()):
            data = data_sets[data_key]
            # plot decay
            n_references = len(data['references'])
            n_plots = 1 + n_references
            fig, ax = plt.subplots(1, n_plots, sharex=True, sharey=True)
            decay = data['decay']
            ax[0].semilogy(decay['time_axis'], decay['decay_histogram'], label=data_key)
            ax[0].semilogy(decay['time_axis'], decay['instrument_response_function'], label='IRF')
            ax[0].semilogy(decay['time_axis'], decay['linearization'], label='linearization')
            ax[0].legend()
            for i, decay_key in enumerate(data['references']):
                decay = data['references'][decay_key]
                ax[i + 1].semilogy(decay['time_axis'], decay['decay_histogram'], label=decay_key)
                ax[i + 1].semilogy(decay['time_axis'], decay['instrument_response_function'], label='IRF')
                ax[i + 1].semilogy(decay['time_axis'], decay['linearization'], label='linearization')
                ax[i + 1].legend()
            plt.show()

    with open(json_file, 'r') as fp:
        fps_json = json.load(fp)
    experiments = fps_json['Experiments']
    data_sets = {}
    for experiment_name in experiments:
        experiment = experiments[experiment_name]
        experiment_file = pathlib.Path(experiment['file'])
        if not experiment_file.is_file():
            parent_path = pathlib.Path(json_file).parent
            print('experiment_file', experiment_file)
            print('parent_path', parent_path)
            experiment_file = parent_path / experiment_file
        info = "Reading experiment file %s" % experiment_file
        skf.logging.info(info)
        if verbose:
            print(info)
        data_sets[experiment_name] = read_tcspc_yaml(experiment_file)
    if plot_data:
        plot_data_sets(data_sets)
    return data_sets
