import pandas as pd
import numpy as np
from .QCBase import VarNames

class Exporter(object):
    """ Export class which writes parsed data to a certain format"""
    valid_formats = ["pdf", "xlsx", "txt", "csv", "dataframe"]
    
    def __init__(self, data=None):
        self.data = data
    
    # for later: add pandas independent functions to export arrays to file
        
    def arrays_to_dframe(self, **kwargs):
        """ Using keyworded arguments, expects arrays  """
        try:
            df = pd.DataFrame(kwargs)
        except ValueError: #if arrays do not have the same length
            d = {}
            for key, value in kwargs.items():
                d[key] = pd.Series(value)
            df = pd.DataFrame(d)
        return df
    
    def ExcitedStateSummary(self, results, fname="es_smry", fmt="csv",
                            ground_state=False):
        """ Exports energy related excited state quantities to file 
        
        Parameters
        ----------
        results : CCParser.ParseContainer
            Parsing container that holds parsed values.
        fname : string
            Filename prefix.
        fmt : string
            Output format ('csv', 'xlsx'/'xls' or 'df' for pandas.DataFrame).
        ground_state : bool
            Whether to include an empty line in the table for the ground state.
        """
        if fmt not in Exporter.valid_formats:
            raise ValueError("File format '{0:}' not recognized or supported!".format(fmt))
        if False in getattr(results, VarNames.has_converged).data:
            raise ValueError("Not converged state detected!")
        d = {}
        # (1) Excitation energies (default minimum)
        #if hasattr(results, VarNames.exc_energy_rel):
        d[VarNames.exc_energy_rel] = getattr(results, VarNames.exc_energy_rel).data
        n_states = len(d[VarNames.exc_energy_rel])
        # (2) Oscillator strengths
        if hasattr(results, VarNames.osc_str):
            d[VarNames.osc_str] = getattr(results, VarNames.osc_str).data
        # (3) Amplitudes
        if hasattr(results, VarNames.amplitudes):
            ampl = getattr(results, VarNames.amplitudes)
            pieces = [a.to_dataframe() for a in ampl]
            key = [x for x in range(1,len(pieces)+1)]
            amp_df = pd.concat(pieces, keys=key, names=["State", "Row ID"])
            
        # prepare MultiIndex (there has to be a better way to do that...)
        arrays = [[x for x in range(1, n_states+1)],
                  [0 for x in range(n_states)]]
        tuples = list(zip(*arrays))# asterisk unpacks
        df1 = pd.DataFrame(d)
        df1.index = pd.MultiIndex.from_tuples(tuples, names=["State", "Row ID"])
        
        df = pd.concat([df1, amp_df], axis=1)
        # add row to MultiIndex, see https://stackoverflow.com/q/24917700
        if ground_state:
            df.loc[(0,0),:] = np.nan
            df.sort_index(level=0, inplace=True)
        
        # EXPORT TO FILE or dataframe
        fout = fname + "." + fmt
        if fmt == "csv":
            df.to_csv(fout, encoding="utf-8")
        elif fmt == ("xlsx" or "xls"):
            writer = pd.ExcelWriter(fout)
            df.to_excel(writer, "Sheet1")
            writer.save()
        elif fmt.lower() == ("dataframe" or "df"):
            return df
        
    def ReducedWeights(self, results, nbsfA, extern=None, fmt="print",
                       fname="AmplAnl", silent=False):
        """ Calculate reduced weights based on fragment information.
        
        The reduced weight for a single excitation :math:`i \\rightarrow a` is defined as       
        :math:`v_{i}^{a} = 0.5\\cdot(c_{i,A}^{2} + c_{a,A}^{2})\\cdot w_{i}^{a}`, with
        c and w being the molecular orbital coefficient and transition weight,
        respectively.
        The MO coefficients from the output first have to be transformed to an
        orthonormal basis.
        
        Parameters
        ----------
        results : CCParser.ParseContainer
            Container object which contains excited state amplitudes
        nbsfA : int
            Number of basis functions on System A (assumes system A comes first!)
        extern : CCParser.ParseContainer
            Optional second container which contains orthonormalisation matrix and/or MO coefficients
        fmt : string
            Output format. Available are "print", "dataframe", "xlsx" or "csv"
        fname : string
            Output file name (basename only).
        silent : bool
            Whether to ignore lengthy printouts.

        """
        # consistency
        has_extern = True if extern != None else False
        if False in getattr(results, VarNames.has_converged).data:
            raise ValueError("Not converged state detected!")