"""
Module to compute the strength of the material using the strength package
"""
import os
from pathlib import Path
import pandas as pd

import numpy as np

from calm.chemistry import ChemistrySizeDistribution
from calm.microstructure import Microstructure 
from strength.yieldstrength import YieldStrength
from premod.parser_generic import parse_complete


# what is the starting point for computing the Yield strength?
#   - microstructure state output
#   - a csv file with only average values
#   - history file


# Get microstructure object
# ============================

# Microstructure object
def get_default_matrix():
    """Create a Microstructure object (calm) from a microstructure dict 
    """
    micro = Microstructure(
        grainsize = 100.E-06,
        delta = 100.E-06, 
        rho_i = 1.E+10,   
        MTaylor = 3.1,   
    )
    return micro


# ============================
# Define precipitates object
# ============================
"""phases = {'beta_pp': {'Al': 5/33, 'Si': 14/33, 'Mg': 14/33},
         'U2': {'Al': 1/3, 'Si': 1/3, 'Mg': 1/3},
         'beta_p': {'Si': 3.33/9.33, 'Mg': 6/9.33},
         'Q_p': {'Al': 6/21, 'Si': 7/21, 'Mg': 6/21, 'Cu': 2/21},
         'L': {'Al': 20/99, 'Si': 33/99, 'Mg': 44/99, 'Cu': 2/99}}"""

def get_default_precipitate_state():
    """
    Default values to test the functions
    """
    metadata = {
        "ndomains": "5",
        "nelements": "6",
        "nphases": "1",
        "elements": ["Al","Mg","Si","Mn","Fe","Cu"],
        "phases": ["MG5SI6_B_DP_MYHR2001"],
        "X0":["0.98155092236025010","0.64466160664951319E-2","0.83177178803558374E-2","0.25927553985026675E-2","0.10202590093591746E-2","0.71729285037052328E-4"],
        #"XSS":[0.99218917962176767,0.12194061391042116E-3,0.39642000370768973E-2,0.26208562319053144E-2,0.10313167930845924E-2,0.72506702255074987E-4],
        "Xd_MG5SI6_B_DP_MYHR2001":["0.0000000000000000","0.58999999999999997","0.40999999999999998","0.0000000000000000","0.0000000000000000","0.0000000000000000"],
    }
    df = pd.DataFrame()
    df["DomainID[-]"] = [1,2,3,4,5]
    df["rpart[m]"] = [0.15000000200271479E-8,0.25000000333785799E-8,0.35000000467300119E-8,0.45000000600814438E-8,0.55000000734328758E-8]
    df["aspect[-]"] = [1.0,1.0,1.0,1.0,1.0]
    df["Nd[1/m3]"] = np.array([0.1,1.0,10.0,1.0,0.01])*100.0E+19
    df["phaseID[-]"] = [1,1,1,1,1]
    return metadata,df

def get_premod_state(filepath):
    """
    Parse a premod state and return the metada and dataframe
    """  
    metadata, df = parse_complete(filepath)

   # get in addition the time and the case name
 
    return metadata, df


def define_chem_premod(metadata:dict,df):
    """
    Define Chemistry object (calm) using data from precipitates state generated by premod.
    Return a ChemistrySizeDistribution object (defined in Calm)
    """
    # hard-coded phase name

    # Create chemistry object
    chem = ChemistrySizeDistribution(
        elements = metadata['elements'],
        X0 = np.asarray(metadata["X0"], dtype=float),
        phases = ['Al_matrix'],
    )
    
    # test to work only with single phase for now
    if int(metadata["nphases"]) > 1:
        raise ValueError("Only single phase microstructure allowed for now in define_chem_premod")

    #assuming only beta''
    #print(metadata)
    # make metadata["phases"] a list if not
    if isinstance(metadata['phases'],str):
        metadata['phases'] = [metadata['phases']]

    for idx, phasename in enumerate(metadata['phases']):
        #print("metadata in defin_chem_premod: ","Xd_"+phasename.strip(),metadata["Xd_"+phasename.strip()])
        chem.add_phase(
            name = phasename.strip(),
            Xp = np.asarray(metadata["Xd_"+phasename.strip()], dtype=float),
            rd = df["rpart[m]"].to_list(),
            Nd = df["Nd[1/m3]"].to_list(),
            rmean = None,
            rstd = None,
            shape = df["aspect[-]"].to_list(),
            distribution = None,
            reduce_volfrac=False,
        )
    
    #print("In define_chem_premod, chem_dict: \n", chem.__dict__)

    #comment = input("wainting for input")
    
    return chem


def define_chem_csv(filepath):
    """
    Define Chemistry object (calm) from a formatted csv file.
    Return a list of ChemistrySizeDistribution object (defined in Calm)
    """
    df_csv = pd.read_csv(filepath,sep=";",comment="#")
    columns = df_csv.columns.to_list()

    """
    X0_ ..., phasenames, MG5SI6_B_DP_MYHR2001_rv_avg (m),MG5SI6_B_DP_MYHR2001_Nv (m-3), Xd_MG5SI6_B_DP_MYHR2001_Al...
    """
    metadata = {"ndomains": 1}
    # extract elements
    elements = get_elements_columns(columns=columns)
    metadata["elements"] = elements
    metadata["nelements"] = len(elements)

    # extract phases
    phases = get_phases_columns(columns=columns)
    metadata["phases"] = phases
    metadata["nphases"] = len(phases)

    list_chems = []

    for irow,row in df_csv.iterrows():
        # Create chemistry object
        chem = ChemistrySizeDistribution(
        elements = metadata['elements'],
        X0 = get_X0_csv(metadata,row),
        phases = ['Al_matrix'],
        )
        for idx, phasename in enumerate(metadata['phases']):
            chem.add_phase(
                name = phasename,
                Xp = get_Xd_csv(metadata,row,phasename),
                rd = df_csv[phasename+"_rv_avg (m)"],
                Nd = df_csv[phasename+"_Nv (m-3)"],
                rmean = None,
                rstd = None,
                shape = df_csv[phasename+"_Shape"],
                distribution = None,
                reduce_volfrac=False,
            )
        list_chems.append(chem)
    
    return list_chems


def get_X0_csv(metadata,row):
    """
    Get the current initial composition from that row in the csv file
    """
    X0 = []
    for element in metadata["elements"]:
        label = "X0_"+element+" (at_frac)"
        try:# label in row.columns:
            X0.append(row[label])
        except:
            raise ValueError("Label: "+label+ " not found in columns names")
    return X0


def get_Xd_csv(metadata,row,phasename):
    """
    Get the current composition of the phase from that row in the csv file
    """
    Xd = []
    for element in metadata["elements"]:
        label = phasename+"_Xd_"+element+" (at_frac)"
        try:# label in row.columns():
            Xd.append(row[label])
        except:
            raise ValueError("Label: "+label+ " not found in columns names")
    return Xd


def get_elements_columns(columns:list):
    """
    Extract from the list the list of elements
    Return: list of elements
    """
    elements = []

    for label in columns:
        if label.startswith("X0_"):
            elements.append(label[3:].split(" ")[0])
    return elements


def get_phases_columns(columns:list):
    """
    Extract the name of the phases from the columns name
    Return list of phases name
    """
    phases = []
    
    for label in columns:
        if label.endswith("_Nv (m-3)"):
            phases.append(label.split("_Nv (m-3)")[0])
    return phases


# Calculate YieldStrength
# ============================
def calculate_YS_premod(metadata:dict,df,micro=None,
                 model='Myhr2015', terms=True):
    """
    Function computing the yield strength of a microstructure for a given model.
    Function expecting instances of entities as dict as input parameters.
    Returns total yield and tau for each contribution term.

    """
    if micro is None:
        micro = get_default_matrix()
    
    chem = define_chem_premod(metadata,df)

    s_tot, tau_terms = calculate_YS(chem,micro,model=model, terms=terms)

    s_terms = {}
    for term, value in tau_terms.items():
        s_terms[term] = value * micro.MTaylor
    return s_tot, s_terms


def calculate_YS_csv(filepath,micro=None,
                 model='Myhr2015', terms=True):
    """
    Function computing the yield strength of a microstructure for a given model.
    Function expecting instances of entities as dict as input parameters.
    Returns total yield and tau for each contribution term.

    """
    if micro is None:
        micro = get_default_matrix()
    
    list_chems = define_chem_csv(filepath)

    list_Ys = []
    for chem in list_chems:
        s_tot, tau_terms = calculate_YS(chem,micro,model=model, terms=terms)
        s_terms = {}
        for term, value in tau_terms.items():
            s_terms[term] = value * micro.MTaylor
        list_Ys.append([s_tot, s_terms])
    return list_Ys


# Calculate YieldStrength
# ============================
def calculate_YS(chem,micro,model='Myhr2015', terms=True):
    """
    Function computing the yield strength of a microstructure for a given model.
    Function expecting instances of entities as dict as input parameters.
    Returns total yield and tau for each contribution term.
    """
    yld = YieldStrength(chem, micro)
    yld.models['ppt'] = model
    strength = yld.yieldstrength(terms = terms)
    s_tot, tau_terms = strength
    return s_tot, tau_terms


def YS_to_HV_AA6xxx(yield_stress):
    """
    Calculation of the Vickers and  hardness based on the yield stress

    The coefficient for Vickers hardness are coming from fitting of 
    some measurements in Ampere
    """
    hardness = {
        "Vickers": 27 + 0.24 * yield_stress,
        "Rockwell B": 50 + 0.32 * yield_stress,
    }

    return hardness


if __name__ == "__main__":

    print("local testing")
    #test_getDistribution()
    #test_defineChem()

    #test_calculate_YS()
