# -*- coding: utf-8 -*-
from rmtk.vulnerability.common import utils
import csv
import numpy as np
import matplotlib.pyplot as plt
from scipy import optimize

def assign_Mw_Tg(input_file, gmrs, Mw_multiplier, damping, filter):

    if filter == 'TRUE':
        Mw, Tg, record_name = [],[],[]
        with open(input_file, 'rU') as f:
            data = csv.reader(f)
            for line in data:
                record_name.append(line[0])
                Mw.append(float(line[1]))
                Tg.append(compute_predominant_period(record_name[-1], gmrs, damping))

        return [record_name, Mw, Tg, Mw_multiplier]
    else:
        return ['None']

def compute_predominant_period(gmr_name, gmrs, damping):

    igmr = gmrs['name'].index(gmr_name)
    T = np.arange(0.01,2,0.01)
    Sv = utils.calculate_iml(igmr, gmrs, 'Sv', T, damping)
    Tg = T[Sv==max(Sv)]

    return Tg[0]

def calculate_fragility_model(PDM,gmrs,gmr_info,IMT,msa,damage_model,T,damping, method):
    
    fragility_model = MSA_post_processing(PDM,gmrs,gmr_info,IMT,msa,damage_model,T,damping, method,0)

    return fragility_model

def calculate_fragility_model_damaged(PDMs,gmrs,gmr_info,IMT,msa,damage_model,T,damping, method):

    fragility_model = {}
    for iDS in range(1,len(PDMs)+1):
        fragility_model['iDS'+str(iDS)] = MSA_post_processing(PDMs['iDS'+str(iDS)],gmrs,np.array(gmr_info['iDS'+str(iDS)]),IMT,msa,damage_model,T,damping, method, iDS)

    return fragility_model

def MSA_post_processing(PDM,gmrs,gmr_info,IMT,msa,damage_model,T,damping, method,initial_damage):
    
    Nrecords = msa['records per bin']
    Nstripes = msa['n. bins']
    no_gmrs = Nrecords*Nstripes
    imls = []
    for index in np.arange(0,no_gmrs,Nrecords):
        igmr = gmr_info[index,0]
        scaling = gmr_info[index,1]
        imls.append(utils.calculate_iml(int(igmr), gmrs, IMT, T, damping)*scaling)
    
#    plt.plot(imls[0],Sds['iDS1'],marker='o', linestyle='None')
#    plt.ylim((0,0.5));
        
    no_assets = np.sum(PDM, axis=1, dtype=float)
    cumPDM = np.fliplr(np.cumsum(np.fliplr(PDM), axis=1))
    cumPDM = np.transpose(cumPDM)
    cumPDM = cumPDM/no_assets

    logmeans, logstddev = [], []
    I = np.argsort(imls, axis=0)
    imls = np.array(imls)
    imls = imls[I]
    if method == 'least squares':
        for iDS in range(initial_damage,len(damage_model['damage_states'])):
            cumPDM[iDS+1] = cumPDM[iDS+1][I]
            solution, _ = optimize.leastsq(utils.residual_lognormal_dist,
                                           [0.1, 0.6],
                                           args=(imls, cumPDM[iDS+1]))
            
            logmeans.append(solution[0])
            logstddev.append(solution[1])

    elif method == 'max likelihood':
        # the imls and the PDM are sorted to be used in the maximum likelihood approach
        # the PDM should be a list of integers in this approach
#        I = np.argsort(imls, axis=0)
#        imls = np.array(imls)
#        imls = imls[I]
        PDM_mle, PoE = [], []
        for iDS in range(initial_damage,len(damage_model['damage_states'])):
            cumPDM[iDS+1] = cumPDM[iDS+1][I]
            temp = cumPDM[iDS+1]*no_assets
            temp2 = temp.astype(int)
            PDM_mle = temp2.tolist()
            solution = utils.mle(PDM_mle,imls,no_assets,1)
            logmeans.append(solution[0])
            logstddev.append(solution[1])
    
    Rsquare = utils.calculate_correlation(logmeans,logstddev,cumPDM[initial_damage:],imls)

    fragility_model = {'damage_states': None, 'logmean': None,'logstddev': None, 'IMT': None}
    fragility_model['damage_states'] = damage_model['damage_states'][initial_damage:]
    fragility_model['logmean'] = logmeans
    fragility_model['logstddev'] = logstddev
    fragility_model['IMT'] = IMT
    fragility_model['Rsquare'] = Rsquare
    fragility_model['PoE points'] = cumPDM[initial_damage+1:]
    fragility_model['imls'] = imls

    return fragility_model

def plot_fragility_model(fragility_model,damage_model,minIML, maxIML):

    for iDS in range(1,len(damage_model['damage_states'])):
        utils.plot_fragility_model(fragility_model['iDS' + str(iDS)],minIML, maxIML)

def save_mean_fragility(fragility_model,damage_model,tax,output_type,output_path,minIML, maxIML):

    for iDS in range(1,len(damage_model['damage_states'])):
        taxonomy = tax+'_iDS'+str(iDS)
        utils.save_mean_fragility(taxonomy, fragility_model['iDS'+str(iDS)], minIML, maxIML, output_type, output_path)

def define_2MSA_parameters(no_bins,no_rec_bin,record_scaled_folder,filter_aftershocks):

    msa = {}; 
    msa['n. bins']=no_bins; 
    msa['records per bin']=no_rec_bin; 
    msa['input folder']=record_scaled_folder; 
    msa['filter'] = filter_aftershocks

    return msa