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

"""
    Methods to perform a GLUE type model calibration. 
    
"""

# Copyright (c) 2014 MTA Centre for Ecological Research
# Distributed under the GNU GPL v3. For full terms see the file LICENSE.

#import numpy.linalg as linalg
#import matplotlib.pyplot as plt
#import re
#import commands
#import pickle
#import struct
#import sys
import os
import numpy
import math
import datetime
import zipfile
#import ResultEvaluator
import wsconf
import bbgc_db
import FileManager
from CalculateFunctions import DerivedOutputs as CaOut
import ReportGenerator
import logging
import logging.config

logging.config.fileConfig('logging.conf')

DAY = 'day'
ANN = 'annual'
MONAVG = 'monthlyaverage'
ANNAVG = 'annualaverage'

class MCalibr:

    def __init__(self,pr_id, mce_id):
        self.pr_id = pr_id
        self.mce_id = mce_id
        self.obs_day_data = []
        self.obs_day_header = []
        self.obs_day_dates = []

        #TODO: the other 4 types (annout, monavg...)



    def calculate_Keenan_MISFIT(self, observed,model_data,squareERR=True,sigma=1.0):
        """ Keenan et al. 2011: eq. 1 alapján
            Normalised Keenan misfit.
            
            :param observed: Observed data.
            :type observed: List.
            :param model_data: Modeled data.
            :type model_data: List. Same length a the observed.
            :param squareERR: Simple or squared Error.
            :type squareERR: Boolean.
            :param sigma:Value of uncertainty. (Not used)
            :type sigma: float
            :returns: Value of misfit (J)

        """

        # Filtering Good values
        nonbadval = [ i for i in range(len(observed)) if observed[i] != None]
        meres = numpy.array([observed[i] for i in nonbadval])
        modell = numpy.array([model_data[i] for i in nonbadval])

        n=len(meres)

        # Diff. of Observed - Model
        if squareERR:
            k=numpy.power((meres-modell)/sigma,2)
        else:
            k=numpy.abs(meres-modell)

        #  Normalised MISFIT - Mean of norm. mean
        return( ((numpy.sum(k))/n) / numpy.mean(meres) )


    def read_observed_csv(self, fname, obs_type=DAY):
        """ Read the observation data from a csv file.
        
            :param fname: Observation data file name.
            :param data: The list, which stores the given observation data.
            :param header: List to store the header of the given observation data.
            :param dates: List to store the dates of observation data.

        """

        if obs_type == DAY:
            header = self.obs_day_header
            data = self.obs_day_data
            dates = self.obs_day_dates
        else:
            return 'Wrong observation type!'

        del header[:]
        del dates[:]
        del data[:]

        first_ln = True
        leap_year = False #nem szoko
        yr = ''
        sep_char = wsconf.csv_separator
        for lin in open(fname):
            if first_ln:
                header.extend(lin.strip().split(sep_char))
                first_ln = False
            elif not first_ln and len(lin.strip()) > 0:
                flds = lin.strip().split(sep_char)
                day_date = flds[0].strip()
                day_date[:4]
                
                #szokoevet megjelolni
                if day_date[4:] == '-2-29' or day_date[4:] == '-02-29':
                    leap_year = True
                    yr = day_date[:4]
                #szokonapot (dec.31) atugrani
                if leap_year and day_date[4:] == '-12-31' and day_date[:4] == yr:
                    leap_year = False
                    yr = ''
                #nem szokoev-dec.31.
                else:
                    dates.append(flds[0].strip())
                    tmp_values = []
                    for v in flds[1:]:
                        try:
                            tmp_values.append(float(v))
                        except ValueError:
                            tmp_values.append(None)
                    data.append(tmp_values)
                    # Ket tombot letrehozni. 1. az adatoknak 2. az uncert. ertekeknek
                    # A datumok egy kulon listbe kerulnek.
                #ha hianyzik dec.31.
                if leap_year and day_date[:4] != yr:
                    leap_year = False
                    yr = ''


    def get_likelihood(self):
        """ Return the likelihood (Stedinger et al. 2008) of the next model run."""
#-------------------------------------------------------------------------------

def get_output_id_from_short_name(short_name):
    """ Get the output variable id from a short name and return in int format.
    
        :param short_name: Short name of the output variable, which contains the (id) string.

    """

    return(int(short_name.split('(')[1].strip(') \n')))

#-------------------------------------------------------------------------------

def get_relevant_model_output(out_first_yr, obs_first_day, out_last_yr, obs_last_day):
    """ Get the first relevant row number and rows number from output data calculated from observation file first day.
    
        :param out_first_yr: first year of output
        :param obs_first_day: first day of observation

    """
    #obs_first_day = '2001-01-11'
    #obs_last_day = '2005-12-31'

    row_nums = {} #dict of numbers of first and last rows
    row_nums['out_start'] = 0
    row_nums['obs_start'] = 0
    obs_first_yr = int(obs_first_day[:4])
    obs_last_yr = int(obs_last_day[:4])

    obs_first_day_date = datetime.datetime.strptime(obs_first_day , '%Y-%m-%d')
    obs_first_yrnum = int(obs_first_day_date.strftime("%j")) -1 #obs first day - day-of-yr, 0-val kezdve
    obs_last_day_date = datetime.datetime.strptime(obs_last_day , '%Y-%m-%d')
    obs_last_yrnum = int(obs_last_day_date.strftime("%j")) -1 #obs last day - day-of-yr, 0-val kezdve
    if obs_last_yrnum == 365:
        obs_last_yrnum = 364

    out_length = (out_last_yr-out_first_yr+1)*365
    obs_length = (obs_last_yr-obs_first_yr)*365 - obs_first_yrnum + (obs_last_yrnum + 1) #sorok szama az obs-ben
    #print 'obs_first_yrnum :' + str(obs_first_yrnum) + ', obs_last_ynum :' + str(obs_last_yrnum)

    if out_first_yr < obs_first_yr: #obs kesobb kezdodik
        first_yr = obs_first_yr
        first_yrnum = obs_first_yrnum #elso vizsgalando ev elso napjanak sorszama, 0-val kezd
        daydiff = (obs_first_yr - out_first_yr)*365 + first_yrnum
        row_nums['out_start'] = daydiff
        out_length = out_length - daydiff
    elif out_first_yr > obs_first_yr: #out kesobb kezdodik
        first_yr = out_first_yr
        yr_diff = out_first_yr - obs_first_yr

        daydiff = yr_diff*365 - obs_first_yrnum
        row_nums['obs_start'] = daydiff
        obs_length = obs_length - daydiff
    else: #ha ugyanabban az evben kezdodik
        first_yr = out_first_yr
        if obs_first_yrnum > 0:
            row_nums['out_start'] = obs_first_yrnum
            out_length = out_length - obs_first_yrnum

    if obs_last_yr < out_last_yr: #obs elobb vegzodik
        last_yr = obs_last_yr #akkor is, ha csonka ev
        daydiff = (out_last_yr - obs_last_yr)*365 + (364-obs_last_yrnum)
        out_length = out_length - daydiff
    elif obs_last_yr > out_last_yr: #out elobb vegzodik
        last_yr = out_last_yr
        yr_diff = obs_last_yr - out_last_yr
        daydiff = (yr_diff-1)*365 + obs_last_yrnum + 1
        obs_length = obs_length - daydiff
    else: #ha ugyanabban az evben vegzodik
        last_yr = out_last_yr
        daydiff = 364-obs_last_yrnum
        out_length = out_length - daydiff

    row_nums['out_length'] = out_length
    row_nums['obs_length'] = obs_length

    return(row_nums)
#-------------------------------------------------------------------------------

def write_result_csvs(project_var_names, glue_results, run_ids, rand_var_names, rand_values, project_id):
    """ Write 2 csv files. 1. Misfit values and likelihood. 2. Randomised input
        variables and likelihood.

        :param project_var_names: Names of examined output variables. (List)
        :param glue_results: Dictionary of the GLUE results. {'run_id':{'partial_llhood':[], 'likelihood':}}
        :param run_ids: Ordered list of run_ids. Same as the order in the random input data.
        :param rand_var_names : List of names of randomised input variables.
        :param rand_values: 2D list of randomised input variables. Rows are the runs.
    """
    sep_char = wsconf.csv_separator
    #f = open('glue_likelihoods.csv', 'w')
    #f.write('run_id' + sep_char + sep_char.join(project_var_names) + sep_char + 'Likelihood\n')
    #for k in glue_results.keys():
    #    f.write(k + sep_char + sep_char.join([str(x) for x in glue_results[k]['misfits'] ]) + sep_char +
    #            str(glue_results[k]['likelihood']) + '\n')
    #f.close()
    
    f = open('randinputs_likelihoods.csv', 'w')
    f.write('run_id' + sep_char + sep_char.join(rand_var_names) + sep_char + 'Likelihood\n')
    for i in range(len(run_ids)):
        rid = run_ids[i]
        f.write(rid + sep_char + sep_char.join([str(x) for x in rand_values[i]]) + sep_char +
                str(glue_results[rid]['likelihood']) + '\n')
    f.close()
    compress = zipfile.ZIP_DEFLATED
    zip_fname = project_id + 'glue_csv.zip'
    zf = zipfile.ZipFile(zip_fname, mode='w')
    #zf.write('glue_likelihoods.csv', compress_type=compress)
    zf.write('randinputs_likelihoods.csv', compress_type=compress)
    zf.close()
    return(zip_fname)

#-------------------------------------------------------------------------------
def create_glue_report(glue_data, glue_results, run_ids, rand_var_names, rand_values, bdb):
    """
    Create a PDF report from Model Calibration.

    :param glue_data: ProjectGlueData object.
    :param glue_results: Dictionary of the GLUE results. {'run_id':{'partial_llhood':[], 'likelihood':}}
    :param run_ids: Ordered list of run_ids. Same as the order in the random input data.
    :param rand_var_names: List of names of randomised input variables.
    :param rand_values: 2D list of randomised input variables. Rows are the runs.
    :param bdb: A bbgc_db.DB object.

    """

    llhood_list = [glue_results[x]['likelihood'] for x in run_ids]
    randin_matrix = numpy.array(rand_values)
    title = ['Biome-BGC GLUE', 'General Likelihood Uncertainty Estimation Report', 'BioVeL 2011 - 2014']
    p_id = str(glue_data.nid)
    mce_id = str(glue_data.mce_nid)
    report_fn = 'GLUE_report_' + mce_id + '_' + p_id + '.pdf'
    rg = ReportGenerator.ReportGenerator(title, report_fn)
    
    rg.add_spacer()
    rg.add_section_title('Project summary')
    rg.add_text('Project title: ' + glue_data.project_data['title'])
    pr_url =  wsconf.drupal_url + 'node/' + p_id
    rg.add_url_line(pr_url, pr_url, 'GLUE project ID: ' + p_id + ' [', ']')
    rg.add_spacer()
    rg.add_text(glue_data.description)
    observ_id = str(glue_data.project_data['field_pr_g_observ'])
    observ_url = wsconf.drupal_url + 'node/' + observ_id
    rg.add_url_line(observ_url, observ_url, 'Observation ID: ' + observ_id + ' [', ']')
    #SELECTED VARIABLES TO CALIBRATE
    rg.add_text('Selected output variables to calibrate:')
    for out_var_name in glue_data.day_vars:
        rg.add_text(' - ' + out_var_name)

    # Time
    start_dt = glue_data.project_data['field_pr_start_time']
    #end_dt = glue_data.project_data['field_pr_end_time']
    now_dt = datetime.datetime.now()
    rg.add_text('Generated: ' + now_dt.strftime('%Y-%m-%d %H:%M'))
    delta_t = now_dt - start_dt
    hours, remainder = divmod(delta_t.seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    rg.add_text('Duration: %s:%s:%s' % (hours, minutes, seconds))

    # MCE
    rg.add_spacer()
    rg.add_text("This analysis is based on the outputs of the " + 
                "following Biome-BGC Monte Carlo Experiment (MCE):")
    rg.add_spacer()
    pr_url = wsconf.drupal_url + 'node/' + mce_id 
    rg.add_url_line(pr_url, pr_url, 'MCE project ID: ' + mce_id + ' [', ']')
    #examined output var names: glue_data.day_vars
    mce_data = bbgc_db.MceIniData(mce_id, bdb)
    rg.add_text(mce_data.project_data['title'])
    rg.add_text('Model type: ' + str(mce_data.model_data['title']))
    rg.add_text('MCE ini ID: ' + str(mce_data.project_data['field_pr_dsin_mceini']))
    rg.add_text('Repetition: ' + str(mce_data.repeat_count))
    rg.add_text('Randomized input variables:')
    for rand_in_name in mce_data.mcefield_list.keys():
        rg.add_text(' - ' + rand_in_name + ' [' + mce_data.mcefield_list[rand_in_name]['minimum'] + ' - '
                    + mce_data.mcefield_list[rand_in_name]['maximum'] + '] ' 
                    + mce_data.mcefield_list[rand_in_name]['dimension'])
    
    # Plots
    for i, rand_var in enumerate(rand_var_names):
        rg.new_page()
        rg.add_text('<b>Likelihood plot</b>')
        rg.add_spacer()
        rg.add_text('randomized parameter: ' + rand_var)
        lplot = ReportGenerator.ScatterPlot('', rand_var,'Likelihood',
                                            randin_matrix[:,i], llhood_list)
        image_name = 'llhood' + str(i) + '.png'
        lplot.create_image(image_name)
        rg.add_figure(image_name,'')
    rg.save_report()
    FileManager.copy_file2proj_dir(report_fn, mce_id)
    bdb.insert_result_file(report_fn, mce_id, 'report', p_id)
    os.system(wsconf.drupal_cache_clear)


#-------------------------------------------------------------------------------

def run_calibration(mce_result, mce_id, cal_id):
    """ Run Model Calibration.
    
        :param mce_result:
        :param mce_pid: The Id of the MCE project.
        :param cal_id: The Id of the ModelCalibration project.

    """

    bdb = bbgc_db.DB()
    # Observation data
    pgd = bbgc_db.ProjectGlueData(cal_id, bdb)
    is_muso = False if int(pgd.model) == wsconf.bbgc411_model_id else True 
    # The actual directory is the MCE project dir!
    #---------------------------------------------
    mcal = MCalibr(cal_id, mce_id)

    mcal.read_observed_csv(pgd.obs_day_file, DAY)
    # Le kell kerni a Proj.GlueD.bol a hasznalando valtozokat es a megfelelo oszlopokat a headerbol
    # A valtozo title mezeje alapjan azonositom a megfelelo header elemet
    #print(pgd.day_vars)
    #print(mcal.obs_day_header)
    # DAY vars
    if len(mcal.obs_day_header) > 1:
        obs_day_data_arr = numpy.array(mcal.obs_day_data)
        obs_var_header = mcal.obs_day_header[1:]
        run_results = {} # key is the job_id + run_id, value is an other dict {'misfits': array of misfits, 'likelihood': lh of the given run}
        run_id_list = []

        #megfelelo napok az outputbol
        first_yr = mce_result.out_data.first_yr
        first_day = mcal.obs_day_dates[0]
        last_yr = first_yr+mce_result.out_data.yr_num-1
        last_day = mcal.obs_day_dates[-1]
        # matching data pointers from model and observation
        dat_pos = get_relevant_model_output(first_yr, first_day, last_yr, last_day)
        obs_start,obs_len,out_start,out_len = (dat_pos['obs_start'],
            dat_pos['obs_length'], dat_pos['out_start'], dat_pos['out_length'])

        for jd in mce_result.valid_job_dirs:
            out_tgz = mce_result.open_one_job_output(jd)
            job_out_dict = mce_result.read_one_job_dayouts(out_tgz) #returns a dict. Key: number (Id) of the run, value: list of output values
            mce_result.close_one_job_output(out_tgz)
            for k in sorted(job_out_dict.keys()):
                run_id = jd + '-' + k
                run_id_list.append(run_id)
                model_data = numpy.array(job_out_dict[k])
                partial_llhood_values = []
                
                for vname in pgd.day_vars:
                    for i in range(len(obs_var_header)):
                        if vname == obs_var_header[i]:           
                            tmp_arr = obs_day_data_arr[:,i]
                            obs_one_var_arr = tmp_arr[obs_start : obs_start+obs_len]
                            if vname in CaOut.calculated_output_names:
                                tmp_list = []
                                for ii in range(out_start, out_start+len(obs_one_var_arr)):
                                    #print(list(model_data[ii,:]))
                                    tmp_list.append(CaOut.get_calculated_output_value(vname,
                                        mce_result.out_data.day_vars,
                                        list(model_data[ii,:]),
                                        is_muso)) #
                                model_col = numpy.array(tmp_list)
                            else:
                                proj_glue_var_id = get_output_id_from_short_name( vname )
                                model_col_index = mce_result.out_data.day_vars.index(
                                    proj_glue_var_id )
                                model_col = model_data[out_start : out_start + len(obs_one_var_arr), model_col_index] #TODO: year!!

                            #misfit_values.append(
                            #    mcal.calculate_Keenan_MISFIT(obs_one_var_arr, model_col, squareERR=pgd.squared_error) )
                            # partial likelihood calculation for each misfit
                            misf = mcal.calculate_Keenan_MISFIT(obs_one_var_arr, model_col, squareERR=pgd.squared_error)

                            sigma = 1.0 # sigma is the uncertainity of the observation, but now is not used
                            partial_llhood = (1.0 / numpy.sqrt(2.0 * math.pi * pow(sigma,2)) ) \
                                             * numpy.exp(-0.5 * misf )
                            partial_llhood_values.append(partial_llhood)
                # Likelihood
                #e=numpy.exp(-0.5 * numpy.mean(misfit_values) )
                #llhood = (1.0 / numpy.sqrt((2.*math.pi*0.16))) * e
                llhood = numpy.prod(numpy.array(partial_llhood_values))

                run_results[run_id] = {'partial_llhood': partial_llhood_values, 'likelihood': llhood}

        #print(run_results)
        in_var_labels = []
        for name in mce_result.randin_names:
            tmp_name = name if name.startswith('field_ds') else 'field_ds_mce_'+name+'1'
            in_var_labels.append(bdb.get_fields_label(tmp_name))

        result_fn = write_result_csvs(pgd.day_vars, run_results, run_id_list,
                                      in_var_labels, mce_result.randin_list, cal_id)
        FileManager.copy_file2proj_dir(result_fn, mce_id)
        bdb.insert_result_file(result_fn, mce_id, 'result', cal_id)
        try:
            create_glue_report(pgd, run_results, run_id_list, 
                               in_var_labels, mce_result.randin_list, bdb)
        except :
            logging.debug('ModelCalibration GLUE report error!')
            import traceback
            logger.debug(traceback.format_exc())
        bdb.close()

    



