#!/usr/bin/env python
"""
Created on Thu Dec 19 14:31:36 2013

Author: Oren Freifeld
Email: freifeld@csail.mit.edu
"""
import numpy as np
from pylab import plt 
from scipy import sparse
from of.utils import *
from cpab.cpaNd import CpaSpace as CpaSpaceNd

from cpab.cpaNd.utils import null
from cpab.cpa2d.utils import *  
from cpab.cpa2d.ConfigPlt import ConfigPlt
 

from cpab.cpa2d.Tessellation import Tessellation
 
class CpaSpace(CpaSpaceNd):
    dim_domain=2
    dim_range=2
    nHomoCoo = dim_domain+1 
    lengthAvee = dim_domain * nHomoCoo
    Ashape =  dim_domain,nHomoCoo

    def __init__(self,XMINS,XMAXS,nCs,
                 zero_v_across_bdry,
                 vol_preserve,warp_around=[False]*2,
                 conformal=False,
                 zero_vals=[],cpa_calcs=None,
                 tess=['II','I'][0],
                 valid_outside=None,
                 only_local=False,
                 cont_constraints_are_separable=None):
        if cont_constraints_are_separable is None:
            raise ObsoleteError("""
            Expected True/False value for cont_constraints_are_separable;
            got None instead""")  
                   
        if tess == 'II' and valid_outside is not None:
            print "tess='II' --> ignoring the value of valid_outside"
        if tess == 'I':
            if valid_outside is None:
                raise ValueError("tess='I' so you must pass valid_outside=True/False" )
            self.valid_outside=valid_outside
                             
        nCx,nCy=map(int,nCs)  

        debug_cont_constraints_are_separable=False
        if cont_constraints_are_separable:
            print 'Check if can actually use separable continuity:'
            if any(zero_v_across_bdry):
                cont_constraints_are_separable=False
                print 'any(zero_v_across_bdry) is True'
            if vol_preserve:
                cont_constraints_are_separable=False  
                print 'vol_preserve is True'
            if nCx!=nCy:
                cont_constraints_are_separable=False
                print 'nCx!=nCy'
            if XMINS[0]!=XMINS[1]:
                cont_constraints_are_separable=False
                print 'XMINS[0]!=XMINS[1]'
            if XMAXS[0]!=XMAXS[1]:
                cont_constraints_are_separable=False
                print 'XMAXS[0]!=XMAXS[1]'
            if not cont_constraints_are_separable:
                debug_cont_constraints_are_separable=False
                print 'so I could not use separable continuity.'
            else:
                print '\nWill use separable continuity.\n'

        super(CpaSpace,self).__init__(XMINS,XMAXS,nCs,
                 zero_v_across_bdry,
                 vol_preserve=vol_preserve,
                 warp_around=warp_around,
                 conformal=conformal,
                 zero_vals=zero_vals,
                 cpa_calcs=cpa_calcs,tess=tess,
                 valid_outside=valid_outside,
                 only_local=only_local,
                 cont_constraints_are_separable=cont_constraints_are_separable) 

        tessellation = Tessellation(nCx,nCy,self.nC,self.XMINS,self.XMAXS,tess=tess)
        self.tessellation=tessellation

        
       
        try:
#                raise FileDoesNotExistError("fake file")
            subspace=Pkl.load(self.filename_subspace,verbose=1)
            B=subspace['B']
            nConstraints=subspace['nConstraints']
            nEdges=subspace['nEdges']
            constraintMat=subspace['constraintMat']
            try:
                cont_constraints_are_separable=subspace['cont_constraints_are_separable']
            except KeyError:
                cont_constraints_are_separable=False
            
        except FileDoesNotExistError: 
            
            nC = self.nC
            verts1,verts2,H,nEdges,nConstraints = self.tessellation.create_verts_and_H(
            dim_range=self.dim_range,valid_outside=valid_outside)
            
            if cont_constraints_are_separable == False or debug_cont_constraints_are_separable:                                                          
                L = create_cont_constraint_mat(H,verts1,verts2,nEdges,nConstraints,nC,
                                               dim_domain=self.dim_domain,
                                               dim_range=self.dim_range)                  
            if cont_constraints_are_separable:
                Lx = create_cont_constraint_mat_separable(H,verts1,verts2,nEdges,nConstraints,
                                                   nC,dim_domain=self.dim_domain,
                                                   dim_range=self.dim_range,tess=tess)
             
            if len(zero_vals): 
                 
                Lzerovals = create_constraint_mat_zerovals(nC,dim_domain=self.dim_domain,
                                                           dim_range=self.dim_range,
                                                           zero_vals=zero_vals)
                L = np.vstack([L,Lzerovals])
                nConstraints += Lzerovals.shape[0]    
                          
            if any(zero_v_across_bdry):
#                    Lbdry = self.tessellation.create_constraint_mat_bdry(
#                                      zero_v_across_bdry=self.zero_v_across_bdry)
#
#                    L = np.vstack([L,Lbdry])
                
                if cont_constraints_are_separable == False or debug_cont_constraints_are_separable:  
                    Lbdry = self.tessellation.create_constraint_mat_bdry(
                                      zero_v_across_bdry=self.zero_v_across_bdry)
    
                    L = np.vstack([L,Lbdry])
                if cont_constraints_are_separable:
                    Lb = self.tessellation.create_constraint_mat_bdry_separable(
                                      zero_v_across_bdry=self.zero_v_across_bdry)
                    raise NotImplementedError(zero_v_across_bdry, cont_constraints_are_separable)                    
                
                nConstraints += Lbdry.shape[0]
            if self.warp_around[0] or self.warp_around[1]:
                raise NotImplementedError
                Lwa = create_constraint_mat_warp_around(cells_verts,
                                                          nC,dim_domain=self.dim_domain)
                L = np.vstack([L,Lwa])
                nConstraints += Lwa.shape[0]
                
            if vol_preserve:
                Lvol = create_constraint_mat_preserve_vol(nC,dim_domain=self.dim_domain)
                L = np.vstack([L,Lvol])
                nConstraints += Lvol.shape[0]
                
            if conformal:
                Lconf = create_constraint_mat_conformal(nC,dim_domain=self.dim_domain,dim_range=self.dim_range)
                L = np.vstack([L,Lconf])
                nConstraints += Lconf.shape[0]                
                
            if self.only_local==False:  
                
                
                if not cont_constraints_are_separable:
                    B=null(L)   
                else: # to solve a nuch smaller SVD and to get a sparser basis                  
                    if vol_preserve or any(zero_v_across_bdry):
                        raise NotImplementedError
                    B1=null(Lx)   
                    # B1.shape is (nC*nHomoCoo)x dim_null_space
                    
                    if debug_cont_constraints_are_separable:
                        B=null(L)
                        if B1.shape[0]!=B.shape[0]/2:
                            raise ValueError(B1.shape,B.shape)
                        if float(B1.shape[1])*self.dim_range != B.shape[1]:
                            raise ValueError(B1.shape,B.shape)
                    _B = np.zeros((B1.shape[0]*2,B1.shape[1]*self.dim_range),B1.dtype)
                    for j in range(B1.shape[1]):
                        Avees = B1[:,j] # length=self.nC*self.nHomoCoo
                        arr=Avees.reshape(self.nC,self.nHomoCoo)
                        for k in range(self.dim_range):
                            arr2=np.hstack([arr if m==k else np.zeros_like(arr) for m in range(self.dim_range)])
                            arr3=arr2.reshape(self.nC,self.lengthAvee)
                            arr4=arr3.flatten()                
                            _B[:,j+k*B1.shape[1]]=arr4
                    if debug_cont_constraints_are_separable:
                        if B.shape != _B.shape:
                            raise ValueError(B.shape,_B.shape)
                    B=_B

                    
                    
            else:
                if tess != 'I':
                    raise NotImplementedError
                B = None
                
            if cont_constraints_are_separable:
                    L=Lx                    
            constraintMat=sparse.csr_matrix(L)  
            
                
            
            Pkl.dump(self.filename_subspace,{'B':B,'cont_constraints_are_separable':cont_constraints_are_separable,
                                             'nConstraints':nConstraints,
                                             'nEdges':nEdges,
                                             'constraintMat':constraintMat},
                                             override=True)


         
           # Since B encodes the null space of, it follows that
           #  np.allclose(L.dot(B),0)==True           
         
        super(CpaSpace,self).__finish_init__(tessellation=tessellation,
                                             constraintMat=constraintMat,
                                             nConstraints=nConstraints,
                                             nInterfaces=nEdges,
                                             B=B,zero_vals=zero_vals) 
                
        self.cont_constraints_are_separable=cont_constraints_are_separable
                   
  
        
       
                     
        self.x_dense = self._calcs.x_dense
        self.x_dense_grid = self._calcs.x_dense_grid   
        
        self.x_dense_img = self._calcs.x_dense_img
        self.x_dense_grid_img = self._calcs.x_dense_grid_img        