#!/usr/bin/env python

'''
utils.py: Utils module to dax-ti-static

__author__     = 'Camilo A. F. Salvador'
__email__      = "camilofs@pm.me"
'''

# -- start

import numpy as np
import pandas as pd

import pymatgen
from pymatgen.core import Lattice, Structure, Molecule, Composition
from pymatgen.core.periodic_table import get_el_sp

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import rc
from matplotlib.ticker import AutoLocator, ScalarFormatter

# Optional settings
plt.style.use('ggplot')
csvpath = 'dax-ti-static.csv'

def load_csv(csvpath: str):
    '''
    Args:
        Path of the *.csv file
    Returns:
        The loaded dataframe (pd.DataFrame)
    Comments:
        Do not remove any columns before importing
    '''

    print('Reading from {}'.format(csvpath))
    df = pd.read_csv(csvpath)
    print ('Database loaded with success; \n')
    return df

def get_at_frac(elem_wts: list):
    '''
    Args:
        A list of tuples with elements and their respective wt_frac   
        e.g. [(Element: string, wt_frac: float)]
    Returns: 
        A list of tuples with elements and their respective at_frac
        e.g. [(Element: string, wt_frac: float)]
    Comments:
        Implements a different syntax than the pymatgen.core.Composition
        to avoid mistaken instances with wt.% instead of at.%
    Usage:
        print ('Conversion from wt -> at%')
        a = get_at_frac([('Ti', 19.0), ('Zr', 21.0), ('Nb', 19.0), ('Mo', 20.0), ('Ta', 21.0)])
        print (a)
        print ('Conversion from at -> wt%') # inherited from pymatgen.Composition
        b = Composition('Ti0.6850Ta0.12Nb0.09Zr0.06O0.015')
        print (b)
        print (10000*100*(a.get_wt_fraction("O")))
    '''

    ids = []
    for a, b in elem_wts: #tuple unpacking
        ids.append(a)
        if isinstance(b, str): # error check
            print ("Error: composition should be a number; \n")
            at_fracs = [] # empty list
            break

    woz = [(b / get_el_sp(a).atomic_mass) for a, b in elem_wts] # weight over Z
    at_fracs = [x / sum(woz) for x in woz] # atomic fractions normalized
    at_fracs = [np.round(float(i), 4) for i in at_fracs] # rounding to :4f
    at_fracs = list(zip(ids, at_fracs)) # make a list
    return at_fracs

def get_moeq(comp: Composition):
    '''
    Args:
        A pymatgen.core.Composition instance   
    Returns: 
        The Mo[eq] parameter (float, 0-100)
    Comment:
        Mo[eq] equation from Koli et al., Metals 8, 506 (2018). 
        doi: http://dx.doi.org/10.3390/met8070506
    '''

    _Mo = comp.get_wt_fraction("Mo")
    _V  = comp.get_wt_fraction("V")
    _W  = comp.get_wt_fraction("W")
    _Nb = comp.get_wt_fraction("Nb")
    _Ta = comp.get_wt_fraction("Ta")
    _Fe = comp.get_wt_fraction("Fe")
    _Cr = comp.get_wt_fraction("Cr")
    _Ni = comp.get_wt_fraction("Ni")
    _Mn = comp.get_wt_fraction("Mn")
    _Co = comp.get_wt_fraction("Co")
    _Al = comp.get_wt_fraction("Al")
    moeq = (_Mo*1.00 + _Nb*0.28 + _Ta*0.22 + _Fe*2.90 + _Cr*1.60 +
            _Ni*1.25 + _Mn*1.70 + _Co*1.70 +  _V*0.67 +  _W*0.44 +
            _Al*(-1))*100
    return moeq

def get_moeq_class(moeq: float):
    '''
    Args:
        The molyybdenum equivalent (Mo[eq]) parameter: float
    Returns: 
        A one in four stability classification: string
    '''
    stability = 'None'
    if moeq > 30.0:
        stability = 'stable'
    elif moeq >= 10:
        stability = 'meta'
    elif moeq >= 5:
        stability = 'near'
    elif moeq >= 0:
        stability = 'rich'
    else:
        stability = 'other'
    return stability

def write_moe_cols(csvpath: str):
    '''
    Args:
        Path of the *.csv file
    Comments:
        Used to add moe and moe_class to the original *.csv
        Allows a quick recalc. of moe if needed
    '''
    
    df = load_csv(csvpath)
    df['formula']     = [x.replace(" ", "") for x in df['formula']] # remove " "
    df['composition'] = [Composition(f) for f in df['formula']]     # pymatgen objs
    df['moe']         = [get_moeq(c) for c in df['composition']]    # calculate moeq
    df['moe_class']   = [get_moeq_class(moe) for moe in df['moe']]  # classification
    df_export = df[['id','eid','formula','R1','R2','R3','F1','F2','F3','P0','P1','YM','YM_err',
                    'YS','YS_err','UTS','UTS_err','DAR','DAR_err','HV','HV_err','moe','moe_class',
                    'ph1','ph2','ph3','condition','comments','composition']]
    df_export.to_csv(csvpath, float_format="%.2f")                  # 2 decimal places
    print('MoE written to {}'.format(csvpath))

def plot_fig1():
    '''
    Comments:
    Plots figure 1 of the article using data from dax-ti-static.csv    
    '''
    df = load_csv(csvpath)
    df = df[['id', 'P0', 'YM', 'YS', 'UTS', 'DAR', 'HV']]
    df.columns = ['id','oc', 'ym', 'ys', 
                  'us', 'ds', 'hv']
    
    columnd = {'oc': 'O content (wppm)', 'ym': 'Young modulus (GPa)',
               'ys': 'Yield strength (MPa)',  'us': 'Ultimate strength (MPa)',
               'ds': 'Deformation (%)', 'hv': 'Hardness (Hv)'}

    ranged =  {'oc': (0,5000), 'ym': (40,160),
               'ys': (0,2000), 'us': (0,2500),
               'ds': (-75,75), 'hv': (0,600)}
    
    xsize = int(2)
    ysize = int(3)
    fig, ax = plt.subplots(xsize, ysize)

    idx = 0
    while idx < (xsize*ysize):
        x = idx % ysize
        y = idx // ysize
        colorn = 'C'+str(idx)

        _col = df.columns[idx+1]
        data = df.iloc[:, [idx+1]]
        
        ax[y, x].hist(data, bins=30, color=colorn)
        ax[y, x].axes.get_yaxis().set_ticks([])
        ax[y, x].set_xlabel(columnd[_col], fontsize=12)
        ax[y, x].set_xlim(ranged[_col])

        print ('\n Data for {}'.format(_col))
        print ('Plot schema: gridpos({}, {}), color = {}'.format(y, x, colorn))
        print (df[_col].describe())
        idx += 1
    
    # changes w/ your mpl.style
    plt.subplots_adjust(left=0.1, bottom=0.15, right=0.9, top=0.9, wspace=0.2, hspace=0.4) 
    plt.show()
    fig.savefig('figure1.png')

def plot_fig2():
    '''
    Comments:
    Plots figure 2 of the article using data from dax-ti-static.csv
    '''
    df = load_csv(csvpath)
    df = df[['id','formula', 'YS', 'DAR', 'moe', 'moe_class']]
    df.columns = ['id','formula', 'ys', 'ds', 'moe', 'class']

    # Filtering outliers
    nan_ds = df['ds'].apply(lambda x: np.isnan(x))
    df = df[nan_ds == False]
    
    nan_ys = df['ys'].apply(lambda x: np.isnan(x))
    df = df[nan_ys == False]

    pos_ds = df['ds'].apply(lambda x: x > 0)
    df = df[pos_ds==True]
    df.loc[df.ds > 50, 'ds'] = 50 # display dar > 50 @ 50

    #classes definition
    classd = {'stable': '$\\beta$-stable', 'meta': '$\\beta$-metastable',
             'near': 'near-$\\beta$', 'rich': '$\\beta$-rich',
             'other':'$\\alpha$+$\\beta$, near-$\\alpha$, c.p.'}
    
    #splitting data into 5 categories (moe_class)
    stable =     df['class'].apply(lambda x: x == 'stable')
    df1 =        df[stable]
    metastable = df['class'].apply(lambda x: x == 'meta')
    df2 =        df[metastable]
    near =       df['class'].apply(lambda x: x == 'near')
    df3 =        df[near]
    rich =       df['class'].apply(lambda x: x == 'rich')
    df4 =        df[rich]
    cp =         df['class'].apply(lambda x: x == 'other')
    df5 =        df[cp]

    # how many samples in each category
    print ('The number of samples in each category:\n')
    print ('Stable: {}'.format(len(df1.axes[0])))
    print ('Meta:   {}'.format(len(df2.axes[0])))
    print ('Near:   {}'.format(len(df3.axes[0])))
    print ('Rich:   {}'.format(len(df4.axes[0])))
    print ('Other:  {}'.format(len(df5.axes[0])))

    # referencial linear fittings
    xs = np.arange(-10, 60) # sample

    l2 = np.polyfit(df2['ds'], df2['ys'], 1)
    l2 = np.poly1d(l2)
    l3 = np.polyfit(df3['ds'], df3['ys'], 1)
    l3 = np.poly1d(l3)
    l4 = np.polyfit(df4['ds'], df4['ys'], 1)
    l4 = np.poly1d(l4)
    l5 = np.polyfit(df5['ds'], df5['ys'], 1)
    l5 = np.poly1d(l5)
        
    fig, ax = plt.subplots()

    ax.scatter(df2['ds'], df2['ys'], label=classd['meta'], color='C1', alpha=0.5)
    ax.plot(xs, l2(xs), '--', color='C1')
    ax.scatter(df3['ds'], df3['ys'], label=classd['near'], color='C2', alpha=0.5)
    ax.plot(xs, l3(xs), '--', color='C2')
    ax.scatter(df4['ds'], df4['ys'], label=classd['rich'], color='C4', alpha=0.8)
    ax.plot(xs, l4(xs), '--', color='C4')
    ax.scatter(df5['ds'], df5['ys'], label=classd['other'],color='C5', alpha=0.8)
    ax.plot(xs, l5(xs), '--', color='C5')
    # only 3 instances of stable alloys; C0 = highest contrast
    ax.scatter(df1['ds'], df1['ys'], color='C0', s=80, alpha=0.8)

    # optional settings
    ax.set_xlim(-5, 55)
    ax.set_yscale('log')
    ax.yaxis.set_major_locator(AutoLocator())
    ax.yaxis.set_major_formatter(ScalarFormatter())
    # ax.minorticks_off()
    ax.set_ylim(200, 2200)

    plt.subplots_adjust(bottom=0.13, top=0.93) 
    plt.legend(prop={'size': 12})
    ax.set_ylabel('Yield strength (MPa)')
    ax.set_xlabel('Deformation at rupture (%)')
    plt.show()
    fig.savefig('figure2.png')


if __name__ == '__main__':
    # std run
    print ('dax-ti utils.py loaded; \n')

# -- end