import pandas as pd
import os
import sys
import warnings as w

from VESIcal import core
from VESIcal import sample_class

def rename_duplicates(df, suffix='-duplicate-'):
    appendents = (suffix + df.groupby(level=0).cumcount().astype(str).replace('0','')).replace(suffix, '')
    return df.set_index(df.index.astype(str) + appendents)

class status_bar(object):
    """Various styles of status bars that display the progress of a calculation
    within a loop
    """
    def __init__():
        pass

    def status_bar(percent, sample_name=None, btext=None, barLen=20):
        """
        Prints an updating status bar to the terminal or jupyter notebook.

        Parameters
        ----------
        percent: float
            Percent value of progress from 0 to 1

        sample_name: string
            Name of the current sample being calculated

        btext: string
            Any extra text to display next to status bar

        barLen: int
            Length of bar to print
        """
        sys.stdout.write("\r")
        sys.stdout.write("[{:<{}}] {:.0f}%".format("=" * int(barLen * percent), barLen, percent * 100))

        sample_string = str(sample_name)
        # Set max number of characters in sample name
        max_name_length = 25
        if len(str(sample_name)) >= max_name_length:
            sample_string = str(sample_name)[0:max_name_length-1] + "..."

        # Write out sample name and trailing spaces to cover contents of previous sample names left over on line
        if sample_name != None:
            sys.stdout.write("  Working on sample " + sample_string + "                            ")
        if btext != None:
            sys.stdout.write(" " + str(btext))
        if percent == 1.0:
            sys.stdout.write("\n")
        sys.stdout.flush()

# ---------- BATCHFILE CLASS --------- #
class BatchFile(object):
    """A batch file with sample names and oxide compositions

    Attributes
    ----------
        filename: str
            Path to the batch file, e.g., "my_file.xlsx". This always needs to be passed, even if the user is passing a pandas DataFrame
            rather than an batch file. If passing a DataFrame, filename should be set to None. File can be excel file (.xlsx) or .csv.

        sheet_name: str
            OPTIONAL. For Excel files. Default value is 0 which gets the first sheet in the batch spreadsheet file. This implements the pandas.
            read_excel() sheet_name parameter. But functionality to read in more than one sheet at a time (e.g., pandas.read_excel(sheet_name=None))
            is not yet imlpemented in VESIcal. From the pandas 1.0.4 documentation:
            
            Available cases:
            - Defaults to 0: 1st sheet as a DataFrame
            - 1: 2nd sheet as a DataFrame
            - "Sheet1": Load sheet with name “Sheet1”

        file_type: str
            OPTIONAL. Default is 'excel', which denotes that passed file has extension .xlsx. Other option is 'csv', which denotes that
            the passed file has extension .csv.

        units: str
            OPTIONAL. Default is 'wtpt_oxides'. String defining whether the oxide composition is given in wt percent
            ("wtpt_oxides", which is the default), mole oxides (mol_oxides) or mole cations (mol_cations).

        default_normalization:     None or str
            The type of normalization to apply to the data by default. One of:
            - None (no normalization)
            - 'standard' (default): Normalizes an input composition to 100%.
            - 'fixedvolatiles': Normalizes major element oxides to 100 wt%, including volatiles.
            The volatile wt% will remain fixed, whilst the other major element oxides are reduced
            proportionally so that the total is 100 wt%.
            - 'additionalvolatiles': Normalises major element oxide wt% to 100%, assuming it is
            volatile-free. If H2O or CO2 are passed to the function, their un-normalized values will
            be retained in addition to the normalized non-volatile oxides, summing to >100%.

        default_units     str
            The type of composition to return by default, one of:
            - wtpt_oxides (default)
            - mol_oxides
            - mol_cations

        label: str
            OPTIONAL. Default is 'Label'. Name of the column within the passed file referring to sample names.

        dataframe: pandas DataFrame
            OPTIONAL. Default is None in which case this argument is ignored. This argument is used when the user wishes to turn
            a pandas DataFrame into an BatchFile object, for example when user data is already in python rather than being imported
            from a file. In this case set `dataframe` equal to the dataframe object being passed in. If using this option, pass
            None to filename.
    """
    def __init__(self, filename, sheet_name=0, file_type='excel', units='wtpt_oxides', label='Label', default_normalization='none', default_units='wtpt_oxides', dataframe=None, **kwargs):
        """Return a BatchFile object whose parameters are defined here."""
        self.units = units
        self.set_default_normalization(default_normalization)
        self.set_default_units(default_units)

        if filename != None:
            file_name, file_extension = os.path.splitext(filename)
            if file_extension == '.xlsx' or file_extension == '.xls':
                file_type = 'excel'
            if file_extension == '.csv':
                file_type = 'csv'

        if isinstance(sheet_name, str) or isinstance(sheet_name, int):
            pass
        else:
            raise core.InputError("If sheet_name is passed, it must be of type str or int. Currently, VESIcal cannot import more than one sheet at a time.")

        if dataframe is not None:
            data = dataframe
            if label is not None:
                data = self.try_set_index(data, label)
        else:
            if file_type == 'excel':
                data = pd.read_excel(filename, sheet_name=sheet_name)
                data = self.try_set_index(data, label)
            elif file_type == 'csv':
                data = pd.read_csv(filename)
                data = self.try_set_index(data, label)
            else:
                raise core.InputError("file_type must be one of \'excel\' or \'csv\'.")

        data = rename_duplicates(data) #handle any duplicated sample names
        data = data.dropna(how='all') #drop any rows that are all NaNs
        data = data.fillna(0) #fill in any missing data with 0's

        if 'model' in kwargs:
            w.warn("You don't need to pass a model here, so it will be ignored. You can specify a model when performing calculations on your dataset \
(e.g., calculate_dissolved_volatiles())",RuntimeWarning,stacklevel=2)

        if 'norm' in kwargs:
            w.warn("We noticed you passed a norm argument here. This does nothing. You can normalize your BatchFile and save it to a new variable name \
after import using normalize(BatchFileObject). See the documentation for more info.",RuntimeWarning,stacklevel=2)

        total_iron_columns = ["FeOt", "FeOT", "FeOtot", "FeOtotal", "FeOstar", "FeO*"]
        for name in total_iron_columns:
            if name in data.columns:
                if 'FeO' in data.columns:
                    for row in data.itertuples():
                        if data.at[row.Index, "FeO"] == 0 and data.at[row.Index, name] > 0:
                            w.warn("Sample " + str(row.Index) + ": " + str(name) + " value of " + str(data.at[row.Index, name]) + " used as FeO. Fe2O3 set to 0.0.",
                                    RuntimeWarning,stacklevel=2)
                            data.at[row.Index, "Fe2O3"] = 0.0
                            data.at[row.Index, "FeO"] = data.at[row.Index, name]
                else:
                    w.warn("Total iron column " + str(name) + " detected. This column will be treated as FeO. If Fe2O3 data are not given, Fe2O3 will be 0.0. \
In future, an option to calcualte FeO/Fe2O3 based on fO2 will be implemented.",RuntimeWarning,stacklevel=2)
                    data['FeO'] = data[name]

        if units == "wtpt_oxides":
            pass
        if units == "mol_oxides":
            data = self._molOxides_to_wtpercentOxides(data)
        if units == "mol_cations":
            data = self._molCations_to_wtpercentOxides(data)

        for oxide in core.oxides:
            if oxide in data.columns:
                pass
            else:
                data[oxide] = 0.0

        self.data = data

    def set_default_normalization(self, default_normalization):
        """ Set the default type of normalization to use with the get_composition() method.

        Parameters
        ----------
        default_normalization:    str
            The type of normalization to apply to the data. One of:
            - 'none' (no normalization)
            - 'standard' (default): Normalizes an input composition to 100%.
            - 'fixedvolatiles': Normalizes major element oxides to 100 wt%, including volatiles.
            The volatile wt% will remain fixed, whilst the other major element oxides are reduced
            proportionally so that the total is 100 wt%.
            - 'additionalvolatiles': Normalises major element oxide wt% to 100%, assuming it is
            volatile-free. If H2O or CO2 are passed to the function, their un-normalized values will
            be retained in addition to the normalized non-volatile oxides, summing to >100%.

        """
        if default_normalization in ['none','standard','fixedvolatiles','additionalvolatiles']:
            self.default_normalization = default_normalization
        else:
            raise core.InputError("The normalization method must be one of 'none', 'standard', 'fixedvolatiles',\
             or 'additionalvolatiles'.")

    def set_default_units(self, default_units):
        """ Set the default units of composition to return when using the get_composition() method.

        Parameters
        ----------
        default_units     str
            The type of composition to return, one of:
            - wtpt_oxides (default)
            - mol_oxides
            - mol_cations
        """
        if default_units in ['wtpt_oxides','mol_oxides','mol_cations']:
            self.default_units = default_units
        else:
            raise core.InputError("The units must be one of 'wtpt_oxides','mol_oxides','mol_cations'.")

    def get_composition(self, species=None, normalization=None, units=None, exclude_volatiles=False, asBatchFile=False):
        """ Returns a pandas DataFrame containing the compositional information for all samples
        in the BatchFile object

        Parameters
        ----------
        species:    NoneType or str
            The name of the oxide or cation to return the concentration of. If NoneType (default) the
            whole composition of each sample will be returned. If an oxide is passed, the value in
            wtpt will be returned unless units is set to 'mol_oxides', even if the default units for the
            sample object are mol_oxides. If an element is passed, the concentration will be returned as
            mol_cations, unless 'mol_singleO' is specified as units, even if the default units for the
            sample object are mol_singleO. Unless normalization is specified in the method call, none
            will be applied.

        normalization:     NoneType or str
            The type of normalization to apply to the data. One of:
            - 'none' (no normalization)
            - 'standard' (default): Normalizes an input composition to 100%.
            - 'fixedvolatiles': Normalizes major element oxides to 100 wt%, including volatiles.
            The volatile wt% will remain fixed, whilst the other major element oxides are reduced
            proportionally so that the total is 100 wt%.
            - 'additionalvolatiles': Normalises major element oxide wt% to 100%, assuming it is
            volatile-free. If H2O or CO2 are passed to the function, their un-normalized values will
            be retained in addition to the normalized non-volatile oxides, summing to >100%.

            If NoneType is passed the default normalization option will be used (self.default_normalization).

        units:     NoneType or str
            The units of composition to return, one of:
            - wtpt_oxides (default)
            - mol_oxides
            - mol_cations
            - mol_singleO
            If NoneType is passed the default units option will be used (self.default_type).

        exclude_volatiles   bool
            If True, volatiles will be excluded from the returned composition, prior to normalization and
            conversion.

        asBatchFile:    bool
            If True, returns a BatchFile object. If False, returns a pandas.DataFrame object.

        Returns
        -------
        pandas.DataFrame or BatchFile object
            All sample information.
        """
        data = self.data.copy()

        # Fetch the default return types if not specified in function call
        if normalization == None and species == None:
            normalization = self.default_normalization
        if units == None and species == None:
            units = self.default_units

        new_compositions = []
        sample_names = []
        for index, row in data.iterrows():
            sample_comp = self.get_sample_composition(index, units=units, asSampleClass=True)
            new_compositions.append(sample_comp.get_composition(species=species, normalization=normalization, units=units, exclude_volatiles=exclude_volatiles))
            sample_names.append(index)
        if isinstance(new_compositions[0], pd.Series):
            return_frame = pd.concat([pd.DataFrame(j) for j in new_compositions], axis=1)
            return_frame = return_frame.transpose()
            return_frame["new_index"] = sample_names
            return_frame = return_frame.set_index("new_index")
            return_frame.index.name = None
        elif isinstance(new_compositions[0], float):
            species_data = {species: new_compositions}
            return_frame = pd.DataFrame(species_data, index=[name for name in sample_names])
        else:
            return_frame = None

        if asBatchFile == False:
            return return_frame
        else:
            return BatchFile(filename=None, dataframe=return_frame, label=None)

    def get_data(self, normalization=None, units=None, asBatchFile=False):
        """
        Returns all data stored in a BatchFile object (both compositional and other data). To return only the
        compositional data, use get_composition().

        Parameters
        ----------
        normalization:     NoneType or str
            The type of normalization to apply to the data. One of:
            - 'none' (no normalization)
            - 'standard' (default): Normalizes an input composition to 100%.
            - 'fixedvolatiles': Normalizes major element oxides to 100 wt%, including volatiles.
            The volatile wt% will remain fixed, whilst the other major element oxides are reduced
            proportionally so that the total is 100 wt%.
            - 'additionalvolatiles': Normalises major element oxide wt% to 100%, assuming it is
            volatile-free. If H2O or CO2 are passed to the function, their un-normalized values will
            be retained in addition to the normalized non-volatile oxides, summing to >100%.

            If NoneType is passed the default normalization option will be used (self.default_normalization).

        units:     NoneType or str
            The units of composition to return, one of:
            - wtpt_oxides (default)
            - mol_oxides
            - mol_cations
            - mol_singleO

            If NoneType is passed the default units option will be used (self.default_type).

        asBatchFile:    bool
            If True, returns a BatchFile object. If False, returns a pandas.DataFrame object.

        Returns
        -------
        pandas.DataFrame or BatchFile object
            All sample information.
        """
        data = self.data.copy()

        # Fetch the default return units if not specified in function call
        if units == None:
            units = self.default_units

        # Grab all compositional data
        compositional_data = self.get_composition(normalization=normalization, units=units)

        # Grab all non-compositional data
        non_compositional_data = data.filter([col for col in data.columns if col not in core.oxides])

        # concatenate both compositional and non-compositional dataframes into one
        return_frame = pd.concat([compositional_data, non_compositional_data], axis=1)

        if asBatchFile == False:
            return return_frame
        else:
            return BatchFile(filename=None, dataframe=return_frame, label=None)

    def get_sample_composition(self, samplename, species=None, normalization=None, units=None, asSampleClass=False):
        """
        Returns oxide composition of a single sample from a user-imported file as a dictionary

        Parameters
        ----------
        samplename: string
            Name of the desired sample

        normalization: NoneType or str
            The type of normalization to apply to the data. One of:
            - 'none' (no normalization)
            - 'standard' (default): Normalizes an input composition to 100%.
            - 'fixedvolatiles': Normalizes major element oxides to 100 wt%, including volatiles.
            The volatile wt% will remain fixed, whilst the other major element oxides are reduced
            proportionally so that the total is 100 wt%.
            - 'additionalvolatiles': Normalises major element oxide wt% to 100%, assuming it is
            volatile-free. If H2O or CO2 are passed to the function, their un-normalized values will
            be retained in addition to the normalized non-volatile oxides, summing to >100%.

            If NoneType is passed the default normalization option will be used (self.default_normalization).

        units:     NoneType or str
            The units of composition to return, one of:
            - wtpt_oxides (default)
            - mol_oxides
            - mol_cations
            - mol_singleO
            If NoneType is passed the default units option will be used (self.default_type).

        asSampleClass:  bool
            If True, the sample composition will be returned as a sample class, with default options. In this case
            any normalization instructions will be ignored.

        Returns
        -------
        dictionary, float, or sample_class.Sample object
            Composition of the sample as oxides
        """
        # Fetch the default return types if not specified in function call
        if normalization == None and species == None:
            normalization = self.default_normalization
        if units == None and species == None:
            units = self.default_units

        # Check that normalization being chosen is one of the possible options
        if normalization in [None, 'none','standard','fixedvolatiles','additionalvolatiles']:
            pass
        else:
            raise core.InputError("The normalization method must be one of 'none', 'standard', 'fixedvolatiles',\
             or 'additionalvolatiles'.")

        data = self.data
        my_sample = pd.DataFrame(data.loc[samplename])
        sample_dict = (my_sample.to_dict()[samplename])
        sample_oxides = {}
        for item, value in sample_dict.items():
            if item in core.oxides:
                sample_oxides.update({item: value})

        _sample = sample_class.Sample(sample_oxides)

        if asSampleClass == True:
            return _sample
        else:
            return_sample = _sample.get_composition(species=species, units=units, normalization=normalization)
            if species == None:
                return dict(return_sample)
            elif isinstance(species, str):
                return return_sample

    def _molOxides_to_wtpercentOxides(self, data):
        for i, row in data.iterrows():
            sample_comp = {}
            for oxide in core.oxides:
                if oxide in data.columns:
                    sample_comp[oxide] = row[oxide]
                else:
                    sample_comp[oxide] = 0.0
            _sample = sample_class.Sample(sample_comp, units='mol_oxides')
            _sample_conv = _sample.get_composition()
            for ox in core.oxides:
                data.loc[i,oxide] = _sample_conv[oxide]
        return data

    def _molCations_to_wtpercentOxides(self, data):
        for i, row in data.iterrows():
            sample_comp = {}
            for cation in core.oxides_to_cations[core.oxides]:
                if cation in data.columns:
                    sample_comp[cation] = row[cation]
                else:
                    sample_comp[cation] = 0.0
            _sample = sample_class.Sample(sample_comp, units='mol_cations')
            _sample_conv = _sample.get_composition()
            for oxide in core.oxides:
                data.loc[i,oxide] = _sample_conv[oxide]
        return data

    def try_set_index(self, dataframe, label):
        """
        Method to handle setting the index column in an BatchFile object. If no column is passed that matches the default index name,
        then this method will attempt to choose the 'best' column that the user might want to serve as an index column.

        Parameters
        ----------
        dataframe: pandas DataFrame

        label: str
            Name of the column within the passed Excel file referring to sample names.
        """
        _dataframe = dataframe.copy()
        try:
            _dataframe = _dataframe.set_index(label)
        except:
            label_found = False
            for col in _dataframe.columns:
                if col in core.oxides:
                    pass
                else:
                    _dataframe = _dataframe.set_index(col)
                    label_found = True
                    w.warn("No Label column given, so column '" + str(col) + "' was chosen for you. To choose your own, set label='<column-name>'.",RuntimeWarning,stacklevel=2)
                    break
            if label_found == False:
                _dataframe.index.name = 'Label'
                w.warn("No Label column given, so one was created for you. To choose your own, set label='<column-name>'.",RuntimeWarning,stacklevel=2)

        return _dataframe

    def preprocess_sample(self,sample):
        """
        Adds 0.0 values to any oxide data not passed.

        Parameters
        ----------
        sample: pandas DataFrame
            self.data composition of samples in wt% oxides

        Returns
        -------
        pandas DataFrame
        """
        for oxide in core.oxides:
            if oxide in self.data.columns:
                pass
            else:
                self.data[oxide] = 0.0

        return sample

    def save_excel(self, filename, calculations, sheet_names=None):
        """
        Saves data calculated by the user in batch processing mode (using the BatchFile class methods) to an organized
        Excel file, with the original user data plus any calculated data.

        Parameters
        ----------
        filename: string
            Name of the file. Extension (.xlsx) should be passed along with the name itself, all in quotes (e.g., 'myfile.xlsx').

        calculations: pandas DataFrame or list of pandas DataFrames
            A single DataFrame or list of DataFrames (e.g., calculated outputs from any of the core BatchFile functions:
            calculate_dissolved_volatiles, calculate_equilibrium_fluid_comp, and calculate_saturation_pressure). If None,
            only the original user data will be saved.

        sheet_names: None, string, or list
            OPTIONAL. Default value is None. Allows user to set the name of the sheet or sheets written to the Excel file.

        Returns
        -------
            Creates and saves an Excel file with data from each calculation saved to its own sheet.
        """
        if isinstance(calculations, list):
            if isinstance(sheet_names, list) or sheet_names is None:
                pass
            else:
                raise core.InputError("If calculations is passed as list, sheet_names must also be list of same length")
        elif calculations == None:
            pass
        else:
            calculations = [calculations]

        with pd.ExcelWriter(filename) as writer:
            self.data.to_excel(writer, 'Original_User_Data')
            if isinstance(calculations, list):
                if sheet_names is None:
                    for n, df in enumerate(calculations):
                        df.to_excel(writer, 'Calc%s' % n)
                elif isinstance(sheet_names, list):
                    pass
                else:
                    sheet_names = [sheet_names]
                if isinstance(sheet_names, list):
                    if len(sheet_names) == len(calculations):
                        pass
                    else:
                        raise core.InputError("calculations and sheet_names must have the same length")

                    for i in range(len(calculations)):
                        if isinstance(sheet_names[i], str):
                            calculations[i].to_excel(writer, sheet_names[i])
                        else:
                            raise core.InputError("if sheet_names is passed, it must be list of strings")
            elif calculations == None:
                pass
        return print("Saved " + str(filename))

    def save_csv(self, filenames, calculations, **kwargs):
        """
        Saves data calculated by the user in batch processing mode to a comma-separated values (csv) file. Mirros the pandas.to_csv()
        method. Any argument that can be passed to pandas.csv() can be passed here. One csv file will be saved for each calculation
        passed.

        Parameters
        ----------
        filenames: string or list of strings
            Name of the file. Extension (.csv) should be passed along with the name itself, all in quotes (e.g., 'myfile.csv'). The number
            of calculations passed must match the number of filenames passed. If passing more than one, should be passed as a list.

        calculations: pandas DataFrame or list of pandas DataFrames
            A single variable or list of variables containing calculated outputs from any of the core BatchFile functions:
            calculate_dissolved_volatiles, calculate_equilibrium_fluid_comp, and calculate_saturation_pressure.

        Returns
        -------
            Creates and saves a CSV file or files with data from each calculation saved to its own file.
        """
        if type(filenames) != list:
            filenames = [filenames]
        if type(calculations) != list:
            calculations = [calculations]
        if len(filenames) != len(calculations):
            raise core.InputError("calculations and filenames must have the same length")

        for i in range(len(filenames)):
            calculations[i].to_csv(filenames[i], **kwargs)
            print ("Saved " + str(filenames[i]))

def from_DataFrame(dataframe, units='wtpt_oxides', label='Label'):
    """
    Transforms any pandas DataFrame object into a VESIcal BatchFile object.

    Parameters
    ----------
    dataframe: pd.DataFrame object
        DataFrame object containing samples and oxide compositions.

    units: str
        OPTIONAL. Default is 'wtpt_oxides'. String defining whether the oxide composition is given in wt percent
        ("wtpt_oxides", which is the default), mole fraction oxides ("mol_oxides"), or mole fraction cations ("mol_cations").

    label: str
        OPTIONAL. Default is 'Label'. Name of the column within the passed file referring to sample names. This
        column will be set as the index column.

    Returns
    -------
    VESIcal.BatchFile object
    """
    return BatchFile(filename=None, dataframe=dataframe, units=units, label=label)
