import pandas as pd
import os
from datetime import date

AbrDict = {'Ala':'ala',
        'Arg':'arg',
        'Asp':'asp',
        'Cys':'cys',
        'Glu':'glu',
        'Gly':'gly',
        'His':'his3',
        'Iso':'ile',
        'Leu':'leu',
        'Lys':'lys',
        'Met':'met',
        'Phe':'phe',
        'Pro':'pro',
        'Ser':'ser',
        'Thr':'thr',
        'Try':'trp',
        'Tyr':'tyr',
        'Val':'val'}
# Automatic generated code with Xcal input:
# def Xcal2iMS2(dataframe, Xcalibur, iMS2):
#     """Convert Xcalibur data to iMS2 format.

#     Parameters
#     ----------
#     dataframe : pandas.DataFrame
#         Dataframe containing Xcalibur data.
#     Xcalibur : str
#         Path to Xcalibur file.
#     iMS2 : str
#         Path to iMS2 file.

#     Returns
#     -------
#     None
#         Writes iMS2 file.
#     """
#     # Read Xcalibur file
#     xcal = pd.read_csv(Xcalibur, sep='\t', header=0)
#     # Get Xcalibur data
#     xcal = xcal[['Scan', 'Precursor', 'Precursor Charge', 'Precursor m/z', 'Precursor Intensity']]
#     # Rename columns
#     xcal.columns = ['Scan', 'Precursor', 'Charge', 'm/z', 'Intensity']
#     # Merge dataframes
#     df = pd.merge(dataframe, xcal, on='Scan')
#     # Write iMS2 file
#     df.to_csv(iMS2, sep='\t', index=False)



def Xcal2iMS(XcalFile, Experiment=None, AAIgnore=None):
    '''Convert Xcalibur data to iMS2Flux input format.
    Parameters
    ----------
    XcalFile : str
        Path to Xcalibur file.
        This file is generated by the Xcalibur software and contains one sheet for each fragment, thus with hundreds of sheets.
        Each sheet contains the data for one fragment, e.g. 'Alanine_232'.
    Experiment : list, optional
        Name of the experiment to be used. The default is None, which corresponds to all experiments.
    AAIgnore : list, optional
        List of amino acids to be ignored. The default is None, which corresponds to no amino acids to be ignored.

    Returns
    -------
    outputfile_name : str
        Writes iMS2Flux tab-separated tsv-file.        
    '''
    Today = date.today().strftime('%y%m%d')
    # Extracting path and filename from XcalFile
    head, tail = os.path.split(XcalFile)

    xls = pd.ExcelFile(XcalFile)
    
    # Reading the first sheet to get the experiment identifiers
    myExp = pd.read_excel(xls, header=4, usecols='A')
    myDat = pd.DataFrame()
    for sheet in xls.sheet_names[:-1]:
        myDat[sheet] = pd.read_excel(xls, sheet, header=4, usecols='O', na_values=['NF'])
    
    # Select defined experiments
    if Experiment is not None:
        myExpList = myExp.iloc[:,0].isin(Experiment)
        myExp = myExp.loc[myExpList,:]
        myDat = myDat.loc[myExpList,:]


    # find columns that contain real data
    UseCols = myDat.notna().any()
    # find rows that contain real data
    UseRows = myDat.iloc[:,0].notna()
    myDat = myDat.loc[:, UseCols] 
    myDat = myDat.loc[UseRows,:]
    # get the fragment names
    myFrags_list = myDat.columns
    myDat = myDat.T
    myDat.reset_index(inplace=True, drop=True)
    # Delete rows that contain NaN in experiment identifiers
    myExp = myExp.loc[UseRows,:]
    myDat.columns = myExp.iloc[:,0]

    # Create new dataframe with fragment names and mass, using AbrDict from above...
    AAId = [AbrDict[Frag[:3]] for Frag in myFrags_list]
    Mass = [Frag[-3:] for Frag in myFrags_list]
    mynew = pd.DataFrame([AAId, Mass]).T
    myRes = pd.concat([mynew, myDat], axis=1)

    # Remove defined amino acid fragments
    if AAIgnore is not None:
        myResList = ~myRes.iloc[:,0].isin(AAIgnore)
        myRes = myRes.loc[myResList,:]
    myRes.reset_index(inplace=True, drop=True)


    # keep only the first name-occurence of each fragment
    repeat = myRes[0].drop_duplicates(keep='first')
    myRes[0] = ''
    myRes.iloc[repeat.index,0] = repeat
    myRes.rename(columns={0:'', 1:''}, inplace=True)

    # Write iMS2Flux file
    outputfile_name = os.path.join(head, 'iMS2Flux_' + tail[:-5] + '_' + Today + '.tsv')
    os.remove(outputfile_name) if os.path.exists(outputfile_name) else None
    with open(outputfile_name, 'a') as file:
        file.write('Response Values (Raw Data)')
        file.write('\n')
        myRes.to_csv(file, sep = '\t', index = False)
    return outputfile_name