# -*- coding: utf-8 -*-

import numpy as np
import scipy.stats as stat
import os
from rmtk.vulnerability.common import utils
from rmtk.vulnerability.mdof_to_sdof import sdof_utils
pi = 3.141592653589793
from rmtk.vulnerability.derivation_fragility.R_mu_T_no_dispersion.dolsek_fajfar.get_spo2ida_parameters import get_spo2ida_parameters

def calculate_fragility(capacity_curves, idealised_capacity, damage_model, MC, Sa_ratios, corner_periods):
    Tc, Td = corner_periods
    allSa, allbTSa, allLR50, allbLR = [],[],[],[]
    no_capacity_curves = len(capacity_curves['Vb'])
    g = 9.81
    w = capacity_curves['weights']
    Tav = capacity_curves['mean_period']
    
    for icc in range(0,no_capacity_curves):
        # Derive median Sa value (median of Sa) of capacity for each Limit State and corresponding overall dispersion std(log(Sa))
        T = capacity_curves['periods'][icc]
        Gamma = capacity_curves['gamma'][icc]
        if len(damage_model['median'])<no_capacity_curves:
            EDPlim = damage_model['median'][0]
            bUthd = damage_model['dispersion'][0]
        else:
            EDPlim = damage_model['median'][icc]
            bUthd = damage_model['dispersion'][icc]
        SPO = idealised_capacity[icc]
        if len(SPO) == 3:
            SPO.insert(2,SPO[1])
            SPO.insert(3,SPO[1])
            SPO.insert(-1,SPO[-1])
            SPO.insert(-1,SPO[-1])

        [ISDvec, Sdvec] = sdof_utils.from_ISD_to_Sd(capacity_curves,icc)
        RDvec = Sdvec*Gamma

        [mc,a,ac,r,mf] = get_spo2ida_parameters(SPO, T, Gamma) # Convert MDoF into SDoF
        [SaT50,bTSa] = DFfragility(T, Gamma, EDPlim, ISDvec, RDvec, SPO, bUthd, mc, r, g, Tc, Td, MC)

        # Converting the Sa(T1) to Sa(Tav), the common IM
        SaTlogmean_av, bTSa_av = np.log(SaT50)*Sa_ratios[icc], np.array(bTSa)*Sa_ratios[icc]
        allSa.append(SaTlogmean_av)
        allbTSa.append(bTSa_av)

    # Combine the fragility of each building in a single lognormal curve with
    # mean = weighted_mean(means) and std = SRSS(weighted_std(means),weighted_mean(stds))
    fragility_model = []
    logmeans, logstddev = [],[]
    for i in range(0,len(EDPlim)):
        SaLS = [ele[i] for ele in allSa]
        StdSaLS = [ele[i] for ele in allbTSa]
        log_m = np.average(SaLS,weights = w)
        log_s = np.sqrt(np.sum(w*(np.power((SaLS-log_m),2)+np.power(StdSaLS,2))))
        logmeans.append(log_m)
        logstddev.append(log_s)

    fragility_model = {'damage_states': None, 'logmean': None, 'logstddev': None,
    'IMT': None}
    fragility_model['damage_states'] = damage_model['damage_states']
    fragility_model['logmean'] = logmeans
    fragility_model['logstddev'] = logstddev
    fragility_model['IMT'] = 'Sa('+str(Tav)+')'

    return fragility_model

def DFfragility(T, Gamma, EDPlim, ISDvec, RDvec, SPO, bUthd, mc, r, g, Tc, Td, MC):
    
#------------------------------------------------------------------------------    
# INPUT
# Cy     : base shear coefficient at yield (presently covered by Gamma and dry)
#          Cy = Say/g = Vy/W, where Say=Sa @ yield, Vy = base shear @ yield
# drlim  : median roof displacement value that defines the fragility. It can result 
#          from any EDP but it should be expressed in terms of the (median)
#          corresponding roof disp (just like we always do in pushovers anyway)
# dry    : roof yield displacement
# buthd  : dispersion (std of log data) characterizing the lognormal
#          distribution of "limit-state roof drift capacity" around drlim
# T      : ESDOF period (sec)
# Gamma  : the participation factor for the roof displacement(>1). Note that for
#          a tall building (or higher-mode influenced one) it is best if you get
#          this value from Say versus droofy estimated from modal response
#          spectrum analysis to include multiple modes. If you use the Gamma of
#          the first one only you will not do well enough (it tends to be
#          lower). See work of Katsanos & Vamva (2014). 
# Tc,Td  : constant accel-constant velocity  and constant velocity-constant
#          displacement corner periods of a Newmark-Hall type spectrum. Default
#          values (roughly taken from Dolsek & Fafjar's (2005) third set of
#          ground motions), are [0.5,1.8].
# mc     : end-of-plastic-plateau capping ductility.
# r      : residual plateau strength divided by yield strength
#
# g      : the value of "g" in units compatible with T and drlim, dry.
#          The default is 9.81m/s2, assuming that dry and drlim are in meters.
# OUTPUT
# Sa50   : the median Sa for the fragility, in units of "g"
# bTSa   : the total dispersion, including capacity and demand dispersions.
#------------------------------------------------------------------------------ 
    if r>1:
        print 'Error: rp should be less than 1'
        os._exit(1)
    if r<0.25:
        print 'Warning: rp should be higher than 0.25, but let us proceed'
    if mc<1:
        print 'Error: mc should be higher than 1'
        os._exit(1)
    if mc>2.5:
        print 'warning: mc should be lower than 2.5, but let us proceed'

    # Relationship between EDP and RD defined by pushover analysis
    if len(ISDvec)>2: #this means that there is a given relationship between ISD e roof displacement (only if dfloor has been input)
        [ISDvec,indy]=np.unique(ISDvec,return_index=True);
        RDvec=RDvec[indy];
        
    drlim=np.interp(EDPlim,ISDvec,RDvec)
        
    dry, du = SPO[0], SPO[-4]
    drlim = np.array(drlim) 
    bUthd = np.array(bUthd)
    if bUthd.any>0 and MC==0:
        print 'error: MC must be different from zero'
        os._exit(1)
    if bUthd.any>0 and MC==0:
        print 'error: MC must be different from zero'
        os._exit(1)
    tdrlim = drlim
    # Ductility level mu for each Limit State
    mlim = np.divide(drlim,dry);
    
    # limiting ductility
    tdrlim[drlim>du] = du
    tmlim = np.divide(tdrlim,dry) # tmlim is limited by du
    print "mu(LS) = ", mlim
    
    if mlim.any()>10: 
        print 'Error: mlim should be less than 10,but let us proceed'
    
    # <codecell>
    
    # For a given period, the IDA model is bilinear (piecewise linear with 2
    # segments)in the post-yield range.
    # In other words it provides a trilinear approximation for the mean IDA (mean mu
    # given R). The point where this change happens is [mc,Rmc].
    # So the mean IDA is defined by 4 points total and 3 segments (incl. elastic):
    # Elastic segment: [0,0] -> [1,1]
    #                  Note though the improvement suggested in D&F (2006) that is
    #                  not adopted here.
    # Capping segment: [1,1] -> [mc,Rmc]
    # Post-capping   : [mc,Rmc] -> [mf,Rmf]
    #                  where the final point is selected by some arbitrary
    #                  Rmf>Rmc, mf=function of Rmf.
    # Thus, to get the mean IDA, I only need mf and Rmf (in addition to mc,Rmc).
    
    # Get the (mc50, Rmc) point
    Tdstar=Td*np.sqrt(2-r)
    R0=1
    m0=1
    # Eq. 2 and 3 in Dolsek and Fajfar (2004)
    if T<=Tc:
        c=0.7*(T/Tc)
    elif T<=Tdstar:
        DT=(T-Tc)/(Tdstar-Tc)
        c=(0.7 + 0.3*DT)
    else:
        c=1
    Rmc = c*(mc-m0) + R0
    
    # from Ruiz-Garcia and Miranda (2007)
    # RGM2007 was fit up to an R=6. Find the largest beta for this period given
    # that R=6 and make sure you do not exceed it.
    R6 = 6
    bthd_max=1.957*(1/5.876+1/(11.749*(T+0.1)))*(1-np.exp(-0.739*(R6-1)));
    
    # find the capping point on the IDA curve.
    if Rmc<=R6:
        bthd_mc=1.957*(1/5.876 + 1/(11.749*(T+0.1)))*(1-np.exp(-0.739*(Rmc-1)))
    else:
        # too large Rmc, cap its dispersion by the max fitted by RGM2007
        bthd_mc=bthd_max
    
    # Get median and fractiles. For lognormal: Median = mean * exp(0.5*sigma^2)
    mc50=mc/np.exp(0.5*np.power(bthd_mc,2))
    mc16=mc50*np.exp(-bthd_mc)
    mc84=mc50*np.exp(bthd_mc)
    
    # <codecell>
    
    # Now set Rmf as twice Rmc and get (mf50,Rmf) point. No real reason why I am using
    # double the value, but since the approximation is linear here, I can
    # interpolate or extrapolate as needed and be fine, so the actual value of Rmf
    # does not matter, as long as the slope Rmf/mf is right.
    Rmf=max(2*Rmc,6)
    R0=Rmc
    m0=mc
    if T<=Tc:
        c=0.7*np.sqrt(r)*np.power((T/Tc),(1/np.sqrt(r)))
    elif T<=Tdstar:
        c=0.7*np.sqrt(r)*(1-DT)+DT
    else:
        c=1
    # This is the mean mu given R at the ficticious Rmf point used to define the
    # post-capping segment.
    mf_mean=(Rmf-R0)/c + m0
    #bthd_mf=1.957*(1/5.876 + 1/(11.749*(T + 0.1)))*(1-np.exp(-0.739*(Rmf - 1)))
    bthd_mf = bthd_max;
    # For lognormal: Median = mean/exp(0.5*sigma^2)
    mf50=mf_mean/np.exp(0.5*np.power(bthd_mf,2))
    mf16=mf50*np.exp(-bthd_mf)
    mf84=mf50*np.exp(bthd_mf)
    
    # <codecell>
    
    R=[0,1,Rmc,Rmf]
    mu16=[0,1,mc16,mf16]
    mu50=[0,1,mc50,mf50]
    mu84=[0,1,mc84,mf84]
    
    # Now we have fractiles and can invert. Use linear extrapolation for mlim values larger than mf.
    c16 = (Rmf-Rmc)/(mf84-mc84)
    c50 = (Rmf-Rmc)/(mf50-mc50)
    c84 = (Rmf-Rmc)/(mf16-mc16)
    cTOT = [c84, c50, c16]
    mcTOT = [mc16, mc50, mc84]
    mfTOT = [mf16, mf50, mf84]
    R16 = np.array([Rmc + c16*(x-mc84) if x>mf84 else np.interp(np.array(x),np.array(mu84),np.array(R)) for x in tmlim])
    R50 = np.array([Rmc + c50*(x-mc50) if x>mf50 else np.interp(np.array(x),np.array(mu50),np.array(R)) for x in tmlim])
    R84 = np.array([Rmc + c84*(x-mc16) if x>mf16 else np.interp(np.array(x),np.array(mu16),np.array(R)) for x in tmlim])
    
    # <codecell>
    
    Say=4*np.power(pi,2)*dry/(g*Gamma*np.power(T,2));    
    SaT50, bTSa = np.zeros_like(mlim),np.zeros_like(mlim)
    
    if bUthd.all() == np.float(0):
        bRSa = 0.5*(np.log(R84)-np.log(R16));
        SaT50 = R50*Say;
        bTSa = bRSa;
    else:
    # Monte Carlo approach
        muf=du/dry
        st = (1./(2.*MC))
        en = (1.-(1./(2.*MC)))
        xp = np.linspace(st,en,MC)
        x=[mu16, mu50, mu84];
        Sai, Sa = [],[]
                
        for i in range(0,len(mlim)):
            Sai.append([])
            Sa.append(np.array([]))
            if bUthd[i]>0:
                EDPsample = stat.lognorm.ppf(xp,bUthd[i],loc=0,scale=EDPlim[i])
                musample = np.interp(EDPsample,ISDvec,RDvec)/dry;
                musample[musample>muf]=muf
                # Estimate R-values of Sa50 and bRSa that correspond to mlim samples
                rMC = [np.interp(musample,ele,R) for ele in x]
                for k in range(0,len(x)):
                    rMC[k][musample>mfTOT[k]] = Rmc+cTOT[k]*(musample[musample>mfTOT[k]]-mcTOT[k])
                    #plt.plot(musample,rMC[k],marker='o')           
                
                allSa50 = [ele*Say for ele in rMC[1]]
                allbSa50 = (np.log(rMC[0])-np.log(rMC[2]))/2
                for j in range(0,MC):
                    # for each potential (equiprobable) sample of mlim, get a sample of Sa-values
                    # as if we actually had N separate IDA curves, distributed according to the Sa50 and bSa50 values extracted above.
                    # In other words,generate N Sa-capacities for N different mlim realizations.
                        if allbSa50[j]>0:
                            realisation = stat.lognorm.ppf(xp,allbSa50[j],loc=0,scale=allSa50[j])
                        else:
                            realisation = np.repeat(allSa50[j],MC)
                        Sai[i].append(realisation)
                        
        for i in range(0,len(mlim)):
            # Find median of sampled Sa-values for those dcroof with bUthd>0, otherwise got to basic RGM method
            if len(Sai[i])>0:
                for j in range(1,len(Sai[i])):
                    Sai[i][j] = np.concatenate((Sai[i][j-1],Sai[i][j]))
                Sa[i] = Sai[i][-1]
                
                SaT50[i] = np.median(Sa[i])
                bTSa[i] = np.std(np.log(Sa[i]))
            else:
                bRSa = 0.5*(np.log(R84)-np.log(R16));
                SaT50[i] = R50[i]*Say;
                bTSa[i] = bRSa[i];
    
# Closed-form Approach, for comparison purposes
    # Go to an R value at 85% of the R50 for a biased estimate of "b", the local
    # slope in log-log of the median IDA.
    Rlim = 0.85*R50
    m50Rlim = np.array([mc50+(Rlim[i]-Rmc)/c50 if mlim[i]>mf50 else np.interp(np.array(Rlim[i]),np.array(R),np.array(mu50)) for i in range(0,len(mlim))])
    b = np.log(m50Rlim)/np.log(Rlim)
    #CR50 = mlim/R50
    Sa50cf=R50*Say;
    bRSacf=0.5*(np.log(R84)-np.log(R16));
    # We are setting here buthd=buEDP for the closed-form solution, i.e., we are
    # assuming proportionality of RD and EDP in the region of interest 
    bUSacf=bUthd/b
    bTSacf = np.sqrt(np.power(bUSacf,2) + np.power(bRSacf,2))
    # if dlim was greater than du it makes little sense to
    # include the dispersion of dlim in the total dispersion. The reason is
    # that we are already at collapse, so changes in dlim have no effect at
    # all. The MC approach takes care of this automatically. Obviously this
    # remains a rough binary approximation (include / don't include) for dlim
    # values close to du.      
    bTSacf[drlim>du] = bRSacf[drlim>du]
    
    print 'medians=',SaT50
    print 'total dispersion=',bTSa
    #plt.close("all")
    return [SaT50, bTSa]
