import numpy as np
import pandas as pd
import os
import typing
import datetime
import warnings

from scipy.optimize import curve_fit
from scipy.optimize import fsolve
from scipy import stats

from . import array as dparray
#from . import csv as dpcsv
from . import integration as integration
from ..file_handling import folder as folder
from ..file_handling import tags as tags


def find_EC_slope(run_dataset: pd.DataFrame, start: float, end: float) -> typing.Tuple[float, float, float]:
    """
    Finds the exponential decay of the EC region for a single dataset. Also returns the intercept and r value of the fit

    Parameters
    ----------
    run_dataset: pd.DataFrame
        dataset of a single DOS run containing at least D/D0, time,
    start: float
        value of D/D0 to start fitting the EC region from
    end: float
        value of D/D0 to end fitting the EC region with

    Returns
    -------
    slope, intercept, r_value: floats
        slope is the slope of the semilog-transformed version of the dataset. It corresponds to the decaying exponential in the linear-linear version of the data
    """

    start_index = dparray.closest_index_for_value(run_dataset, "D/D0", start)
    end_index = dparray.closest_index_for_value(run_dataset, "D/D0", end)
    dataset_EC = run_dataset[start_index:end_index]
    log_diameter = np.log(dataset_EC['D/D0'])
    slope, intercept, r_value, p_value, std_err = stats.linregress(dataset_EC['time (s)'],log_diameter)
    return slope, intercept, r_value, std_err

def annotate_summary_df(fitting_results_list: list, header_params: dict) -> pd.DataFrame:
    """
    Do we want to bring other columns with us like ion, polymer identity, etc? How to code that?

    Parameters
    ----------
    fitting_results_list: list
        generated by find_EC_slope
    header_params : dict
        Contains the information profiled in the samplename

    Returns
    -------
    lambdaE_df : pd.DataFrame
        dataframe containing lambdaE relaxation time for each run from the input df
    """
    # TODO: make fittings_results_list entry in docstring clearer

    ## TODO: address issue raised in docstring

    lambdaE_df = pd.DataFrame(fitting_results_list)
    # Constructs column headers for the summary dataframe
    constant_fitting_header = ("-b","Intercept", "R","slope standard error", "run","Dtc/D0")
    df_header = {}
    keys_list = header_params.keys()
    keys_list = list(keys_list)
    for i in range(0, len(header_params)):
        df_header[i] = keys_list[i]
    for i in range(len(header_params), (len(header_params) + len(constant_fitting_header))):
        df_header[i] = constant_fitting_header[i - len(header_params)]

    lambdaE_df = lambdaE_df.rename(columns=df_header)
    try:
        lambdaE_df['Lambda E (s)'] = -1/(3*lambdaE_df['-b'])
        lambdaE_df['Lambda E (ms)'] = lambdaE_df['Lambda E (s)'] * 1000
        lambdaE_df['R^2'] = (lambdaE_df['R']) ** 2
        lambdaE_df['Lambda E standard error (s)'] = lambdaE_df["slope standard error"]/(3*(lambdaE_df['-b']**2))
        lambdaE_df['Lambda E standard error (ms)'] = lambdaE_df['Lambda E standard error (s)'] * 1000
        lambdaE_df = lambdaE_df.drop(["-b", "Intercept", "R", "slope standard error", 'Lambda E standard error (s)', "Lambda E (s)", ], axis=1)
    except KeyError:
        #print("Error in fitting csv")
        pass

    return lambdaE_df

def make_summary_dataframe(df: pd.DataFrame, sampleinfo_format: str, optional_settings: dict = {}) -> pd.DataFrame:
    """
    Condenses a DOS run into an extensional relaxation time by fitting the EC region (t > tc) to a decaying exponential

    Parameters
    ----------
    df : pd.DataFrame
        Contains D/D0, time, t - tc, strain rate, D(tc)/D0, etc for multiple runs and samples
        generated from data_processing.csv.generate_df
    sampleinfo_format : str
        the format of the sampleinfo section of the filename
        separated by the deliminator specified by sample_split
    optional_settings: dict
        A dictionary of optional settings.

    Optional Settings and Defaults
    ------------------------------
    fitting_bounds: 2 element list of floats
        [start, end]
        The D/D0 to bound the start and end of fitting of EC region.
        Default is [0.1, 0.045].

    Returns
    -------
    summary_df : pd.DataFrame
        dataframe containing lambdaE (relaxation time) and R(t_c)/R_0 for each run from the input df, along with their sample info
    """

    # Initalizes parameters and empty list
    settings = integration.set_defaults(optional_settings)
    fitting_bounds = settings["fitting_bounds"]
    verbose = settings["verbose"]
    start = fitting_bounds[0]
    end = fitting_bounds[1]

    fitting_results_list = []


    samples = df["sample"].unique()
    for sample in samples:
        # Grabs sample info from "sample" field
        header_params = tags.parse_fname(sample,"sampleinfo",sampleinfo_format,optional_settings)
        # Selects individual sample from df
        sample_dataset = df[(df["sample"] == sample)]
        run_values = sample_dataset['run'].unique()
        for run in run_values:
            if verbose:
                print("Fitting Sample: " + str(sample) + " Run: " + str(run))
            run_dataset = sample_dataset[(sample_dataset['run'] == run)]
            run_dataset = run_dataset.reset_index(drop=True)
            Dtc_D0 = run_dataset.loc[0, "Dtc/D0"]
            if Dtc_D0 < start:
                start = Dtc_D0
            try:
                fitting_results_temp =  [*header_params.values(), *find_EC_slope(run_dataset, start, end),run, Dtc_D0]
                fitting_results_list.append(fitting_results_temp)
            except ValueError:
                print("Error in fitting csv of Sample: " + str(sample) + " Run: " + str(run))
    #### TODO: Clean up the dataframe column names ###
    summary_df = annotate_summary_df(fitting_results_list, header_params)
    return summary_df

def derivative_EC_fit(Dtc_D0: float, lambdaE: float, time: float, tc: float) -> float:
    """
    Calculates the derivative of the elasto-capillary region.

    D(t)/D0 = D(tc)/D0 * (exp(-(t - tc)/(3*LambdaE)))
    D'(t)/D0 = (-1/(3*LambdaE)) * D(tc)/D0 * (exp(-(t - tc)/(3*LambdaE)))

    Parameters
    ----------
    Dtc_D0: float
        The normalized diameter at which the transition to EC behavior occurs
    LambdaE: float
        The relaxation time of the polymer solution
    time: float
        the moment in time to evaluate the derivative
    tc: float
        the critical time for the experiment, this number is purely emperical
        (it depends entirely on when in the process the video starts)

    Returns
    -------
    D'(t)/D0: float

    """

    return Dtc_D0*(-1/(3*lambdaE))*np.exp(-(time - tc)/(3*lambdaE))

def calculate_elongational_visc(df: pd.DataFrame, summary_df: pd.DataFrame, optional_settings: dict = {}) -> pd.DataFrame:
    """
    Calculates the quantity (elongational viscosity / surface tension) for each moment in the DOS dataset.

    Parameters
    ----------
    df : pd.DataFrame
        Contains D/D0, time, t - tc, strain rate, D(tc)/D0, etc for multiple runs and samples
        generated from data_processing.csv.generate_df
    summary_df : pd.DataFrame
        Contains relaxation time, D(t_c)/D0, and sample info for all runs and samples
        generated from data_procescing.fitting.make_summary_dataframe
    optional_settings: dict
        A dictionary of optional settings.

    Optional Settings and Defaults
    ------------------------------
    needle_diameter_mm: float
        The needle outer diameter in millimeters.
        Default is 0.7176 mm (22G needle).


    Returns
    ------
    dataset_w_visc : pd.DataFrame
        Dataset with the additional values calculated by this function: strain and (elongational viscosity / surface tension)
    """

    # calculate timestep from the first two rows of data
    # (could also just read the value of time at position 1?)
    #timestep_s = df.loc[1, "time (s)"] - df.loc[0, "time (s)"] #units of seconds

    #get mean relaxation time and D(tc)/D0 from summary_df for current sample
    #Calculate strain and elongational viscosity / surface tension
    #Append them to the dataframe

    settings = integration.set_defaults(optional_settings)
    needle_diameter_mm = settings["needle_diameter_mm"]

    df_w_visc_list = []
    mean_summary_df = summary_df.groupby("sample").mean()
    for sample in summary_df["sample"].unique():
        #subset_sample = summary_df[summary_df["sample"] == str(sample)]
        Dtc_D0_mean = mean_summary_df.loc[sample, "Dtc/D0"]
        lambdaE_mean = mean_summary_df.loc[sample, "Lambda E (ms)"]
        #calculate elongational viscosity from original DOS data and mean values of Dtc/D0 and lambdaE
        for run in df[df["sample"] == sample]["run"].unique():
            dataset = df[df["sample"]==str(sample)]
            dataset = dataset[dataset["run"] == run]
            dataset['strain'] = -2*np.log(dataset['D/D0'])
            dataset=dataset.reset_index(drop=True)

            for value in range(0,len(dataset)):
                t_minus_tc = dataset.at[value, "t - tc (s)"]
                t_minus_tc_ms = t_minus_tc*1000
                # NOTE: Because we prefer the needle diameter to give the length scale,
                # if you use radius, you need to multiply a factor of 2 into the denominator of the following equation
                e_visc_sigma = -1*(1/((derivative_EC_fit(Dtc_D0_mean,lambdaE_mean,t_minus_tc_ms,0))*needle_diameter_mm))
                # e visc / surface tension [=] -1/D'(t) = [1/(mm/ms)] = [1/(m/s)] = [s/m]
                dataset.at[value, "(e visc / surface tension) (s/m)"] = e_visc_sigma
            df_w_visc_list.append(dataset)
    dataset_w_visc = pd.concat(df_w_visc_list)
    return dataset_w_visc

def save_summary_df(summary_df: pd.DataFrame, save_location: typing.Union[str, bytes, os.PathLike], optional_settings: dict = {}):
    """
    Saves the summary dataset from a large processed batch of videos.

    The summary dataset includes only the fluid properties, not the 'raw' diameter vs time data.

    Parameters
    ----------
    summary_df : pd.DataFrame
        Contains LambdaE, D(tc)/D0, and sample info for multiple runs and samples
    save_location: path-like
        path to folder in which to save the csv
    filename: string
        Currently not used, but this will be the name to give to the summary dataset csv
    optional_settings: dict
        A dictionary of optional settings.

    Optional Settings and Defaults
    ------------------------------
    skip_existing: bool
        Determines the behavior when a file already appears exists
        when a function would generate it. True to skip any existing files.
        False to overwrite (or delete and then write, where overwriting would
        generate an error).
        Default is True.
    verbose: bool
        Determines whether processing functions print statements as they
        progress through major steps. True to see print statements, False to
        hide non-errors/warnings.
        Default is False.
    summary_filename: string
        The base filename (no extension) for saving the summary csvs. If not
        provided, will be generated automatically based on the current date
        and time.
        Default is "" to trigger automatic generation.

    Returns
    -------
    filename_string: string
        Filename at which the summary dataframe was saved.
    Saves file to disk.
    """

    settings = integration.set_defaults(optional_settings)
    skip_existing = settings["skip_existing"]
    filename = settings["summary_filename"]
    verbose = settings["verbose"]

    date_and_time = datetime.datetime.now()
    # No colons or periods in filename string.
    if filename == '':
        date_time_string = str(date_and_time.date()) + '_'+str(date_and_time.hour)+'-'+str(date_and_time.minute)+'-'+str(date_and_time.second)
        filename_string = date_time_string + '_DOS-summary.csv'
    else:
        if '_DOS-summary.csv' in filename:
            filename_string = filename
        elif '.csv' in filename:
            filename_string = filename.replace('.csv','') + '_DOS-summary.csv'
        else:
            filename_string = filename + '_DOS-summary.csv'
    full_save_path = os.path.join(save_location,filename_string)

    # If file with given or generated filename does not exist, saves, otherwise,
    # checks skip_existing to see if generates a warning or removes and rewrites
    # the file.
    saved = False
    if os.path.exists(full_save_path):
        if not skip_existing:
            os.remove(full_save_path)
            summary_df.to_csv(full_save_path)
            saved = True
        else:
            warnings.warn("Summary file with given filename already exists and optional_settings skip_existing is True by default. Summary file was NOT saved. Please rerun the function with a new summary_filename specified in optional_settings or with no filename specified.", UserWarning)
    else:
        summary_df.to_csv(full_save_path)
        saved = True
    if verbose:
        if saved:
            print("Summary file saved successfully with name " + filename_string)

    return filename_string

def save_processed_df(df: pd.DataFrame, save_location: typing.Union[str, bytes, os.PathLike], optional_settings: dict = {}):
    """
    Saves the processed dataset from a batch of videos.

    The summary dataset includes the fluid properties, the 'raw' diameter vs time data, and the sampleinfo from the filename.

    Parameters
    ----------
    df : pd.DataFrame
        Raw dataset with calculated columns for strain rate, strain, etc, and parsed sampleinfo columns.
    save_location: path-like
        Path to folder in which to save the csv.
    optional_settings: dict
        A dictionary of optional settings.

    Optional Settings and Defaults
    ------------------------------
    skip_existing: bool
        Determines the behavior when a file already appears exists
        when a function would generate it. True to skip any existing files.
        False to overwrite (or delete and then write, where overwriting would
        generate an error).
        Default is True.
    verbose: bool
        Determines whether processing functions print statements as they
        progress through major steps. True to see print statements, False to
        hide non-errors/warnings.
        Default is False.
    summary_filename: string
        The base filename (no extension) for saving the summary csvs. If not
        provided, will be generated automatically based on the current date
        and time.
        Default is "" to trigger automatic generation.

    Returns
    -------
    filename_string: string
        Filename at which the annotated dataframe was saved.
    Saves file to disk.
    """

    settings = integration.set_defaults(optional_settings)
    skip_existing = settings["skip_existing"]
    filename = settings["summary_filename"]
    verbose = settings["verbose"]

    date_and_time = datetime.datetime.now()
    # No colons or periods in filename string.
    if filename == '':
        date_time_string = str(date_and_time.date()) + '_'+str(date_and_time.hour)+'-'+str(date_and_time.minute)+'-'+str(date_and_time.second)
        filename_string = date_time_string + '_DOS-annotated.csv'
    else:
        if '_DOS-annotated.csv' in filename:
            filename_string = filename
        elif '.csv' in filename:
            filename_string = filename.replace('.csv','') + '_DOS-annotated.csv'
        else:
            filename_string = filename + '_DOS-annotated.csv'
    full_save_path = os.path.join(save_location,filename_string)

    # If file with given or generated filename does not exist, save, otherwise,
    # check skip_existing to see if generate a warning or remove and rewrite
    # the file.
    saved = False
    if os.path.exists(full_save_path):
        if not skip_existing:
            os.remove(full_save_path)
            df.to_csv(full_save_path)
            saved = True
        else:
            warnings.warn("Annotated file with given filename already exists and optional_settings skip_existing is True by default. Annotated file was NOT saved. Please rerun the function with a new summary_filename specified in optional_settings or with no filename specified.", UserWarning)
    else:
        df.to_csv(full_save_path)
        saved = True
    if verbose:
        if saved:
            print("Annotated file saved successfully with name " + filename_string)

    return filename_string
