#!/usr/bin/env python
# ct_statevector_tools.py

"""
Author : peters 

Revision History:
File created on 28 Jul 2010.

"""

import os
import sys
sys.path.append(os.getcwd())

import logging
import numpy as np
from da.baseclasses.statevector import StateVector, EnsembleMember

import da.tools.io4 as io

identifier = 'CarbonTracker Statevector '
version = '0.0'

################### Begin Class CO2C13StateVector ###################

class CO2C13StateVector(StateVector):
    """ This is a StateVector object for CarbonTracker. It has a private method to make new ensemble members """

    def get_covariance(self, date, dacycle):
        """ Make a new ensemble from specified matrices, the attribute lag refers to the position in the state vector. 
            Note that lag=1 means an index of 0 in python, hence the notation lag-1 in the indexing below.
            The argument is thus referring to the lagged state vector as [1,2,3,4,5,..., nlag]
        """    
        try:
            import matplotlib.pyplot as plt
        except:
            pass

        # Get the needed matrices from the specified covariance files

        file_ocn_cov = dacycle.dasystem['ocn.covariance'] 
        file_bio_cov = dacycle.dasystem['bio.covariance'] 

        # replace YYYY.MM in the ocean covariance file string

        file_ocn_cov = file_ocn_cov.replace('2000.01', date.strftime('%Y.%m'))

        for fil in [file_ocn_cov, file_bio_cov]:
            if not os.path.exists(fil):
                msg = "Cannot find the specified file %s" % fil
                logging.error(msg)
                raise IOError, msg
            else:
                logging.info("Using covariance file: %s" % fil)

        f_ocn = io.ct_read(file_ocn_cov, 'read')
        f_bio = io.ct_read(file_bio_cov, 'read')

        cov_ocn = f_ocn.get_variable('CORMAT')
        if f_bio.variables.has_key('covariance'):
            cov_bio = f_bio.get_variable('covariance')  # newly created CTDAS covariance files
            cov_epsbio = f_bio.get_variable('covariance')
            cov_epsbio = cov_bio[:,:]
        else:
            cov_bio = f_bio.get_variable('qprior')  # old CarbonTracker covariance files
            cov_epsbio = f_bio.get_variable('qprior')
        f_ocn.close()
        f_bio.close()
        print 'cov_bio',cov_bio.shape,cov_bio.sum(),cov_bio.max(),cov_bio.min()      
        print 'cov_epsbio',cov_epsbio.shape,cov_epsbio.sum(),cov_epsbio.max(),cov_epsbio.min()
       
        logging.debug("Succesfully closed files after retrieving prior covariance matrices")

        # Once we have the matrices, we can start to make the full covariance matrix, and then decompose it

        fullcov = np.zeros((self.nparams, self.nparams), float)
        offdiag = np.zeros((209, 209), float)
        correlation=.0
        #np.fill_diagonal(offdiag,0.08)
        offdiag=np.sqrt(cov_bio)*np.sqrt(cov_epsbio/16.)*correlation
        nocn = cov_ocn.shape[0]
        nbio = cov_bio.shape[0]

        fullcov[0:nbio, 0:nbio] = cov_bio
        fullcov[nbio:nbio + nocn, nbio:nbio + nocn] = cov_ocn
        fullcov[nbio+nocn, nbio+nocn] = 1.e-10
        fullcov[nbio+nocn+1:nbio+nocn+nbio+1,nbio+nocn+1:nbio+nocn+nbio+1]  = cov_epsbio/16.
        fullcov[nbio+nocn+1:nbio+nocn+nbio+1,0:nbio] = offdiag
        fullcov[0:nbio,nbio+nocn+1:nbio+nocn+nbio+1] = offdiag

        print 'fullcov full',fullcov.diagonal()[:],fullcov.diagonal()[:].shape        
        print 'fullcov bio',fullcov.diagonal()[:nbio],fullcov.diagonal()[:nbio].shape
        print 'fullcov ocn',fullcov.diagonal()[nbio:nbio+nocn],fullcov.diagonal()[nbio:nbio+nocn].shape
        print 'fullcov eps',fullcov.diagonal()[nbio+nocn+1:nbio+nocn+nbio+1],fullcov.diagonal()[nbio+nocn+1:nbio+nocn+nbio+1].shape
        print 'fullcov eps2',fullcov.diagonal()[239:],fullcov.diagonal()[239:].shape
        logging.debug('fullcov bio %s, %s' %(fullcov.diagonal()[:nbio],fullcov.diagonal()[:nbio].shape))
        logging.debug('fullcov eps3 %s, %s' %(fullcov.diagonal()[240:],fullcov.diagonal()[240:].shape))
        print 'fullcov eps4',fullcov.diagonal()[241:],fullcov.diagonal()[241:].shape


        try:
            plt.imshow(fullcov)
            plt.colorbar()
            plt.savefig('fullcovariancematrix_eps.png')
            plt.close('all')
            logging.debug("Covariance matrix visualized for inspection")
        except:
            pass

        print 'fullcov',fullcov.shape,fullcov.sum()
        return fullcov

    def grid2vector_eps(self,griddata=None, method = 'avg'):

        methods =  ['avg','sum','minval']
        if method not in methods:
            msg = "To put data from a map into the statevector, please specify the method to use (%s)" % methods ; logging.error(msg)
            raise ValueError

        result=np.zeros((self.nparams,),float)
 
        for k,v in self.griddict.iteritems():
            if k <210:
                if method == "avg":
                    result[k+240-1]=griddata.take(v).mean()
                elif method == "sum" :
                    result[k+240-1]=griddata.take(v).sum()
                elif method == "minval" :
                    result[k+240-1]=griddata.take(v).min()
        return result # Note that the result is returned, but not yet placed in the member.ParameterValues attrtibute!


    def vector2grid_eps(self,vectordata=None):
        """
            Map vector elements to a map or vice cersa

           :param vectordata: a vector dataset to use in case `reverse = False`. This dataset is mapped onto a 1x1 grid and must be of length `nparams`
           :rtype: ndarray: an array of size (360,180,)

           This method makes use of a dictionary that links every parameter number [1,...,nparams] to a series of gridindices. These
           indices specify a location on a 360x180 array, stretched into a vector using `array.flat`. There are multiple ways of calling
           this method::

               griddedarray = self.VectorToGrid(vectordata=ParameterValues) # simply puts the ParameterValues onto a (180,360,) array

           .. note:: This method uses a DaSystem object that must be initialzied with a proper parameter map. See :class:`~da.baseclasses.dasystem` for details

        """
        #result=np.zeros(self.gridmap.shape,float)
        result=np.ones(self.gridmap.shape,float)

        for k,v in self.griddict.iteritems():
            if vectordata.shape[0] == 209 and k < 210:
                result.put(v,vectordata[k-1])
            if vectordata.shape[0] == 240 or vectordata.shape[0] == 449:
                result.put(v,vectordata[k-1])
        return result


    def vector2tc_eps(self,vectordata,cov=False):
        """
            project Vector onto TransCom regions

           :param vectordata: a vector dataset to use, must be of length `nparams`
           :param cov: a Boolean to specify whether the input dataset is a vector (mean), or a matrix (covariance)
           :rtype: ndarray: an array of size (23,) (cov:F) or of size (23,23,) (cov:T)
        """

        M = self.tcmatrix
        if cov:
            return np.dot(np.transpose(M[:209,:]), np.dot(vectordata[240:],M[:209,:]) )
        else:
            return np.dot(vectordata[240:].squeeze(),M[:209,:])

    def state_to_grid_eps(self, fluxvector = None, lag=1):
        """
            Transforms the StateVector information (mean + covariance) to a 1x1 degree grid.

            :param: fluxvector: a vector of length (nparams,) that holds the fluxes associated with each parameter in the StateVector
            :param: lag: the lag at which to evaluate the StateVector
            :rtype: a tuple of two arrays (gridmean,gridvariance) with dimensions (180,360,)

            If the attribute `fluxvector` is not passed, the function will return the mean parameter value and its variance on a 1x1 map.

            ..note:: Although we can return the variance information for each gridbox, the covariance information contained in the original ensemble is lost when mapping to 1x1 degree!

        """





        if fluxvector == None:
            fluxvector = np.ones(self.nparams)

        ensemble = self.ensemble_members[lag - 1]
        ensemblemean = ensemble[0].param_values

        # First transform the mean
        gridmean = self.vector2grid_eps(vectordata=ensemblemean[240:] * fluxvector[240:])

        # And now the covariance, first create covariance matrix (!), and then multiply
        deviations = np.array([mem.param_values[240:] * fluxvector[240:] - ensemblemean[240:] for mem in ensemble])
        ensemble = []
        for mem in deviations:
            ensemble.append(self.vector2grid_eps(mem))

        return (gridmean, np.array(ensemble))




    def write_members_to_file(self, lag, outdir):
        """ 
           :param: lag: Which lag step of the filter to write, must lie in range [1,...,nlag]
           :rtype: None

           Write ensemble member information to a NetCDF file for later use. The standard output filename is 
           *parameters.DDD.nc* where *DDD* is the number of the ensemble member. Standard output file location 
           is the `dir.input` of the DaCycle object. In principle the output file will have only two datasets inside 
           called `parametervalues` which is of dimensions `nparameters` and `parametermap` which is of dimensions (180,360). 
           This dataset can be read and used by a :class:`~da.baseclasses.observationoperator.ObservationOperator` object. 

           .. note:: if more, or other information is needed to complete the sampling of the ObservationOperator you
                     can simply inherit from the StateVector baseclass and overwrite this write_members_to_file function.

        """

        # These import statements caused a crash in netCDF4 on MacOSX. No problems on Jet though. Solution was
        # to do the import already at the start of the module, not just in this method.

        #import da.tools.io as io
        #import da.tools.io4 as io

      
        members = self.ensemble_members[lag]
    

        for mem in members:
            filename = os.path.join(outdir, 'parameters.%03d.nc' % mem.membernumber)
            ncf = io.CT_CDF(filename, method='create')
            dimparams = ncf.add_params_dim(self.nparams-209)
            dimgrid = ncf.add_latlon_dim()

            data = mem.param_values
            data=data[:240]

            savedict = io.std_savedict.copy()
            savedict['name'] = "parametervalues"
            savedict['long_name'] = "parameter_values_for_member_%d" % mem.membernumber
            savedict['units'] = "unitless"
            savedict['dims'] = dimparams
            savedict['values'] = data
            savedict['comment'] = 'These are parameter values to use for member %d' % mem.membernumber
            ncf.add_data(savedict)

            griddata = self.vector2grid(vectordata=data)

            savedict = io.std_savedict.copy()
            savedict['name'] = "parametermap"
            savedict['long_name'] = "parametermap_for_member_%d" % mem.membernumber
            savedict['units'] = "unitless"
            savedict['dims'] = dimgrid
            savedict['values'] = griddata.tolist()
            savedict['comment'] = 'These are gridded parameter values to use for member %d' % mem.membernumber
            ncf.add_data(savedict)

            dimparams2       = ncf.add_dim('nparameters_epsbio',209)

            data            =  mem.param_values
            data=data[240:]
            savedict                = io.std_savedict.copy()
            savedict['name']        = "parametervalues_epsbio"
            savedict['long_name']   = "parameter_values_for_member_%d"%mem.membernumber
            savedict['units']       = "unitless"
            savedict['dims']        = dimparams2
            savedict['values']      = data
            savedict['comment']     = 'These are parameter values to use for member %d'%mem.membernumber
            ncf.add_data(savedict)

            griddata        = self.vector2grid_eps(vectordata=data)

            savedict                = io.std_savedict.copy()
            savedict['name']        = "parametermap_epsbio"
            savedict['long_name']   = "parametermap_for_member_%d"%mem.membernumber
            savedict['units']       = "unitless"
            savedict['dims']        = dimgrid
            savedict['values']      = griddata.tolist()
            savedict['comment']     = 'These are gridded parameter values to use for member %d'%mem.membernumber
            ncf.add_data(savedict)

            ncf.close()

            logging.debug('Successfully wrote data from ensemble member %d to file (%s) ' % (mem.membernumber, filename))







    def read_from_legacy_file(self, filename, qual='opt'):
        """ 
        :param filename: the full filename for the input NetCDF file
        :param qual: a string indicating whether to read the 'prior' or 'opt'(imized) StateVector from file
        :rtype: None

        Read the StateVector information from a NetCDF file and put in a StateVector object
        In principle the input file will have only one four datasets inside 
        called:
            * `meanstate_prior`, dimensions [nlag, nparamaters]
            * `ensemblestate_prior`, dimensions [nlag,nmembers, nparameters]
            * `meanstate_opt`, dimensions [nlag, nparamaters]
            * `ensemblestate_opt`, dimensions [nlag,nmembers, nparameters]

        This NetCDF information can be written to file using 
        :meth:`~da.baseclasses.statevector.StateVector.write_to_file`

        """
        

        f = io.ct_read(filename, 'read')

        for n in range(self.nlag):
            if qual == 'opt':
                meanstate = f.get_variable('xac_%02d' % (n + 1))
                EnsembleMembers = f.get_variable('adX_%02d' % (n + 1))

            elif qual == 'prior':
                meanstate = f.get_variable('xpc_%02d' % (n + 1))
                EnsembleMembers = f.get_variable('pdX_%02d' % (n + 1))

            if not self.ensemble_members[n] == []:
                self.ensemble_members[n] = []
                logging.warning('Existing ensemble for lag=%d was removed to make place for newly read data' % (n + 1))

            for m in range(self.nmembers):
                newmember = EnsembleMember(m)
                newmember.param_values = EnsembleMembers[m, :].flatten() + meanstate  # add the mean to the deviations to hold the full parameter values
                self.ensemble_members[n].append(newmember)

        f.close()

        logging.info('Successfully read the State Vector from file (%s) ' % filename)


    def make_new_ensemble(self, lag, covariancematrix=None):
        """ 
        :param lag: an integer indicating the time step in the lag order
        :param covariancematrix: a matrix to draw random values from
        :rtype: None
    
        Make a new ensemble, the attribute lag refers to the position in the state vector. 
        Note that lag=1 means an index of 0 in python, hence the notation lag-1 in the indexing below.
        The argument is thus referring to the lagged state vector as [1,2,3,4,5,..., nlag]

        The optional covariance object to be passed holds a matrix of dimensions [nparams, nparams] which is
        used to draw ensemblemembers from. If this argument is not passed it will ne substituted with an 
        identity matrix of the same dimensions.

        """    
        import matplotlib.pyplot as plt

        if covariancematrix == None: 
            covariancematrix = np.identity(self.nparams)

        # Make a cholesky decomposition of the covariance matrix


        C = np.linalg.cholesky(covariancematrix)

        try:
            plt.imshow(C)
            plt.colorbar()
            plt.savefig('fullcholesky_eps2.png')
            plt.close('all')
            logging.debug("cholesky matrix visualized for inspection")
        except:
            pass
        logging.debug('Cholesky decomposition has succeeded ')
        #logging.info('Appr. degrees of freedom in covariance matrix is %s' % (int(dof)))


        # Create mean values 

        newmean = np.ones(self.nparams, float) # standard value for a new time step is 1.0

        # If this is not the start of the filter, average previous two optimized steps into the mix

        if lag == self.nlag - 1 and self.nlag >= 3:
            newmean += self.ensemble_members[lag - 1][0].param_values + \
                                           self.ensemble_members[lag - 2][0].param_values 
            newmean = newmean / 3.0

        # Create the first ensemble member with a deviation of 0.0 and add to list

        newmember = EnsembleMember(0)
        newmember.param_values = newmean.flatten()  # no deviations
        self.ensemble_members[lag].append(newmember)

        # Create members 1:nmembers and add to ensemble_members list
        rand=[]
        rand2=[]
        for member in range(1, self.nmembers):
            rand.append(np.random.randn(240))

        for member in range(1, self.nmembers):
            rand2.append(np.random.randn(209))

        rand = np.array(rand)
        rand2 = np.array(rand2)
        rand=np.hstack((rand,rand2))

        for member in range(1, self.nmembers):
            rands = rand[member-1,:]
            logging.debug( 'rands std %s' %rand[member-1,:].std())

            newmember = EnsembleMember(member)
            a=np.zeros((449),)
            a[:]=np.dot(C[:,:], rands) + newmean
            newmember.param_values = a
            self.ensemble_members[lag].append(newmember)
            if member <10:
                logging.debug('a[0], %s, %s, %s, %s' %(member,a[0],newmean[0],C[0,0]))
                logging.debug('a_eps[0], %s, %s, %s, %s' %(member,a[240],newmean[240],C[240,240]))

        logging.debug('%d new ensemble members were added to the state vector # %d' % (self.nmembers, (lag + 1)))



################### End Class CO2C13StateVector ###################


if __name__ == "__main__":
    from da.tools.initexit import StartLogger 
    from da.tools.pipeline import start_job

    sys.path.append(os.getcwd())

    opts = ['-v']
    args = {'rc':'da.rc', 'logfile':'da_initexit.log', 'jobrcfilename':'test.rc'}

    StartLogger()

    DaCycle = start_job(opts, args)

    DaCycle.Initialize()

    StateVector = CtStateVector()

    StateVector.Initialize()

    for n in range(dims[0]):
        cov = StateVector.get_covariance()
        dummy = StateVector.make_new_ensemble(n + 1, cov)

    StateVector.propagate()

    savedir = DaCycle['dir.output']
    filename = os.path.join(savedir, 'savestate.nc')

    dummy = StateVector.WriteToFile()

    StateVector.ReadFromFile(filename)

