# -*- coding: utf-8 -*-
import os
import numpy
import math
import matplotlib.pyplot as plt
from rmtk.vulnerability.common import utils

def generate_capacity_curves(Sa_mean,Sa_cov,Sd_mean,Sd_cov,type_dist,no_cc,Sa_corr,Sd_corr,SaSd_corr,truncation):
#Creates a set of randomly-generated capacity curves, based on a mean curve.
      
    no_points = len(Sa_mean)
    mean, covariance = calculate_variables(Sa_mean,Sd_mean,Sa_cov,Sd_cov,Sa_corr,Sd_corr,SaSd_corr)
    zero_col = numpy.zeros((no_cc,1))
    Sa = []
    Sd = []
    generated_cc = 0
        
    while generated_cc < no_cc:
        sample = numpy.random.multivariate_normal(mean,covariance,1)
        if validate_sample(sample[0],truncation,Sa_mean,Sa_cov,Sd_mean,Sd_cov):
            Sa.append(sample[0][0:no_points])
            Sd.append(sample[0][no_points:])
            generated_cc = generated_cc +1
    
    Sa = numpy.array(Sa)
    Sd = numpy.array(Sd)
    Sa_corr = numpy.corrcoef((Sa[:,-1],Sa[:,-2]))
    Sd_corr = numpy.corrcoef((Sd[:,-1],Sd[:,-2]))
    Sa_Sd_corr = numpy.corrcoef((Sa[:,-1],Sd[:,-1]))
    
    print 'Actual correlation in Sa: '+str(Sa_corr[0][1])
    print 'Actual correlation in Sd: '+str(Sd_corr[0][1])
    print 'Actual correlation between Sa and Sd: '+str(Sa_Sd_corr[0][1])

    Sa = numpy.hstack((zero_col,Sa))
    Sd = numpy.hstack((zero_col,Sd))
    
    capacity_curves = {'type': None, 'periods': None, 'heights': None, 
        'gamma': None, 'Sdy': None, 'Say': None, 'Sd': None, 'Sa': None}

    capacity_curves['type'] = 'Sd-Sa'
    capacity_curves['Sd'] = Sd
    capacity_curves['Sa'] = Sa    
    
    return capacity_curves
    
def validate_sample(sample,truncation,Sa_mean,Sa_cov,Sd_mean,Sd_cov):
#This function performs a series of sanity checks in order to validate the randomly generated
#sample
    
    result = True
    no_points = len(Sa_mean)
    Sa = sample[0:no_points]
    Sd = sample[no_points:]
    delta = Sd[1:]-Sd[0:-1]
    Sa_sigma = numpy.multiply(Sa_mean,Sa_cov)
    Sd_sigma = numpy.multiply(Sd_mean,Sd_cov)
    Sa_epsilon = abs(numpy.divide((Sa-Sa_mean),Sa_sigma))
    Sd_epsilon = abs(numpy.divide((Sd-Sd_mean),Sd_sigma))

    if min(sample) < 0:
        result = False

    if min(delta) <= 0:
        result = False
        
    if max(Sa_epsilon) > truncation:
        result = False
        
    if max(Sd_epsilon) > truncation:
        result = False
    
    return result
    
def calculate_variables(Sa_mean,Sd_mean,Sa_cov,Sd_cov,Sa_corr,Sd_corr,SaSd_corr):
#This function returns the mean values of Sa and Sd values and calculates the covariance matrix  
   
    mean = []
    for value in Sa_mean:
        mean.append(value)
    for value in Sd_mean:
        mean.append(value)  

    cov = []
    for value in Sa_cov:
        cov.append(value)
    for value in Sd_cov:
        cov.append(value)
        
    no_points = len(mean)
    covariance = numpy.zeros((no_points,no_points))        
    
    for irow in range(no_points):
        for icol in range(no_points):
            sigma1 = mean[irow]*cov[irow] 
            sigma2 = mean[icol]*cov[icol]
            if irow == icol:
                covariance[irow][icol] = sigma1*sigma2
            else:
                if irow < no_points/2 and icol < no_points/2:
                    corr =  Sa_corr
                elif irow < no_points/2 and icol >= no_points/2:
                    corr =  SaSd_corr
                elif irow >= no_points/2 and icol < no_points/2:
                    corr =  SaSd_corr
                elif irow >= no_points/2 and icol >= no_points/2:
                    corr =  Sd_corr
                
                covariance[irow][icol] = sigma1*sigma2*corr
    return mean, covariance
    
def add_information(capacity_curves,attribute,type,data):
#Adds additional information to the capacity_curves dictionary (like storey heights, period,participation
#factors,etc.)
    
    no_capacity_curves = len(capacity_curves['Sa'])
    
    if attribute == 'heights' or attribute == 'periods' or attribute == 'gamma':           
        if type == 'value':
            values = []
            for icc in range(no_capacity_curves):
                values.append(data)
            capacity_curves[attribute] = values
        elif type == 'vector':
            capacity_curves[attribute] = data
        elif type == 'calculate' and attribute == 'periods': 
            periods = []
            for icc in range(no_capacity_curves):
                Sd = capacity_curves['Sd'][icc][data] 
                Sa = capacity_curves['Sa'][icc][data] 
                periods.append(2*math.pi*math.sqrt(Sd/(Sa*9.81)))
            capacity_curves[attribute] = periods
            
    elif attribute == 'yielding point':
        Sdy = []
        Say = []
        for icc in range(no_capacity_curves):
            Sdy.append(capacity_curves['Sd'][icc][data])
            Say.append(capacity_curves['Sa'][icc][data])
        capacity_curves['Sdy'] = Sdy 
        capacity_curves['Say'] = Say 
    else:
        print attribute + ' is not a recognized attribute. No information was added.'
        
    return capacity_curves