import aquila_borg as borg
import time
import numpy as np

# Emulator imports
import torch
torch.cuda.empty_cache() 

from networks import StyledVNet, StyledVNet_distill

# JAX set-up
import jax
from functools import partial
import jax.numpy as jnp
from jax import vjp
from jax.config import config as jax_config
jax_config.update("jax_enable_x64", True)

from utils import myprint
import utils

class Emulator(borg.forward.BaseForwardModel): 

    def __init__(self, box, prev_module, NN, Om, upan, use_float64, q, requires_grad, cuda_avail, debug=False):
        super().__init__(box, box)
        
        # Set-up and cosmo
        self.box = box
        self.q = q 
        self.Om = Om
        self.debug = debug
        
        # LPT and EMU
        self.prev_module = prev_module
        self.NN = NN
        
        # Since we won't return an adjoint for the density, but for positions, we need to do:
        self.prev_module.accumulateAdjoint(True)
        
        # Settings 
        self.requires_grad = requires_grad
        self.use_pad_and_NN = upan
        self.use_float64 = use_float64
        self.cuda_avail = cuda_avail
        if self.use_float64:
            self.dtype = np.float64
        else:
            self.dtype = np.float32 
        
    def requires_grad_true(self):
        self.requires_grad = True
       
    def requires_grad_false(self):
        self.requires_grad = False
   
    # IO "preferences"
    def getPreferredInput(self):
        return borg.forward.PREFERRED_REAL

    def getPreferredOutput(self):
        return borg.forward.PREFERRED_REAL

    # Forward part
    def forwardModel_v2_impl(self, input_array):
        global requires_grad

        if self.debug:
            myprint(f'Requires grad = {self.requires_grad}')
        
        init_time = time.time()
        
        # Step 0 - Extract particle positions
        if self.debug:
            start_time = time.time()
        pos = np.zeros((self.prev_module.getNumberOfParticles(), 3)) #output shape: (N^3, 3)
        self.prev_module.getParticlePositions(pos)
        if self.debug:
            myprint("Step 0 of forward pass took %s seconds" % (time.time() - start_time))
        
        # Step 1 - find displacements
        if self.debug:
            start_time = time.time()
        disp = pos - self.q
        if self.debug:
            myprint("Step 1 of forward pass took %s seconds" % (time.time() - start_time))
        
        # Step 2 - correct for particles that moved over the periodic boundary
        if self.debug:
            start_time = time.time()
        disp_temp = utils.correct_displacement_over_periodic_boundaries(disp,L=self.box.L[0],max_disp_1d=self.box.L[0]//2)
        if self.debug:
            myprint("Step 2 of forward pass took %s seconds" % (time.time() - start_time))
        
        # Step 3 - reshaping initial pos and displacement
        if self.debug:
            start_time = time.time()
        # not sure why order='C' is working here... not sure if it matters... could change it below
        q_reshaped = np.reshape(self.q.T, (3,self.box.N[0],self.box.N[0],self.box.N[0]), order='C') #output shape: (3, N, N, N)
        dis_in = np.reshape(disp_temp.T, (3,self.box.N[0],self.box.N[0],self.box.N[0]), order='C') #output shape: (3, N, N, N)
        if self.debug:
            self.dis_in = dis_in
            myprint("Step 3 of forward pass took %s seconds" % (time.time() - start_time))
        
        # Step 4 - normalize
        if self.debug:
            start_time = time.time()
        utils.dis(dis_in)
        if self.debug:
            myprint("Step 4 of forward pass took %s seconds" % (time.time() - start_time))
        
        if self.use_pad_and_NN:
            if self.debug:
                myprint('Using padding and NN.')
            
            if not self.use_float64:
                dis_in = dis_in.astype(np.float32)
                
            # Step 5 - padding to (3,N+48*2,N+48*2,N+48*2)
            if self.debug:
                start_time = time.time()
            
            # for jax adjoint
            if self.debug:
                myprint(f'in_pad shape = {dis_in.shape}')
            dis_in_padded, self.ag_pad = vjp(self.padding, dis_in) #output shape: (3, N+96, N+96, N+96)
            if self.debug:
                myprint(f'out_pad shape = {np.shape(np.asarray(dis_in_padded))}')
                myprint("Step 5 of forward pass took %s seconds" % (time.time() - start_time))
            
            # Step 6 - turn into a pytorch tensor (unsquueze because batch = 1)
            if self.debug:
                start_time = time.time()
            if self.use_float64:
                self.x = torch.unsqueeze(torch.tensor(np.asarray(dis_in_padded),dtype=torch.float64, requires_grad=self.requires_grad),dim=0) #output shape: (1, 3, N+96, N+96, N+96)
            else:
                self.x = torch.unsqueeze(torch.tensor(np.asarray(dis_in_padded),dtype=torch.float32, requires_grad=self.requires_grad),dim=0) #output shape: (1, 3, N+96, N+96, N+96)
            if self.cuda_avail:
                self.x = self.x.cuda() 
            if self.debug:
                myprint("Step 6 of forward pass took %s seconds" % (time.time() - start_time))
            
            # Step 7 - Pipe through emulator  
            if self.debug:
                start_time = time.time()
            if self.requires_grad:
                self.y = self.NN(self.x,self.Om) #output shape: (1, 3, N, N, N)
            else:
                with torch.no_grad():
                    self.y = self.NN(self.x,self.Om) #output shape: (1, 3, N, N, N)
            if self.debug:
                myprint("Step 7 of forward pass took %s seconds" % (time.time() - start_time))
            
            # Step 8 - N-body sim displacement
            if self.debug:
                start_time = time.time()
            dis_out = torch.squeeze(self.y).detach().cpu().numpy()
            
        else:
            myprint('Skipping padding and NN.')
            dis_out = dis_in
        if self.debug:
            myprint("Step 8 of forward pass took %s seconds" % (time.time() - start_time))
        
        # Step 9 - undo the normalization
        if self.debug:
            start_time = time.time()
        utils.dis(dis_out,undo=True)
        if self.debug:
            self.dis_out = dis_out
            myprint("Step 9 of forward pass took %s seconds" % (time.time() - start_time))
        
        # Step 10 - convert displacement into positions
        if self.debug:
            start_time = time.time()
        pos = dis_out + q_reshaped
        if self.debug:
            myprint("Step 10 of forward pass took %s seconds" % (time.time() - start_time))
        
        # Step 11 - make sure everything within the box
        if self.debug:
            start_time = time.time()
        pos[pos>=self.box.L[0]] -= self.box.L[0]
        pos[pos<0] += self.box.L[0]
        if self.debug:
            myprint("Step 11 of forward pass took %s seconds" % (time.time() - start_time))
        
        # Step 12 - reshape positions
        if self.debug:
            start_time = time.time()
        self.pos_out = pos.reshape(3,self.box.N[0]**3,order='C').T #output shape: (N^3, 3)
        if self.debug:
            myprint("Step 12 of forward pass took %s seconds" % (time.time() - start_time))
        
        # Step 13 - CIC
        if self.debug:
            start_time = time.time()
        #self.dens_out = cic_analytical(self.pos_out[:, 0], self.pos_out[:, 1], self.pos_out[:, 2], *self.box.N + self.box.L)     
        self.dens_out, _ = vjp(self.jax_cic, self.pos_out[:, 0], self.pos_out[:, 1], self.pos_out[:, 2])
        
        if self.debug:
            myprint("Step 13 of forward pass took %s seconds" % (time.time() - start_time))
            myprint('Finished forward pass of emu in %s seconds' % (time.time() - init_time))
        
    def getParticlePositions(self,output_array):
        output_array[:] = self.pos_out
        
    def getDensityFinal_impl(self, output_array):
        output_array[:] =  self.dens_out
        
    # Adjoint part
    def adjointModel_v2_impl(self, input_ag):

        init_time = time.time()
        # reverse step 13 (CIC)
        if self.debug:
            start_time = time.time()
        ag = np.asarray(self.cic_analytical_grad(input_ag,self.pos_out[:,0],self.pos_out[:,1],self.pos_out[:,2],128,128,128,250,250,250))
        ag = np.copy(ag) 
        if self.debug:
            myprint("Reverse step 13 took %s seconds" % (time.time() - start_time))
        
        # reverse step 11
        if self.debug:
            start_time = time.time()
        ag = np.reshape(ag.T, (3,*self.box.N), order='C') #changed recently
        if self.debug:
            myprint("Reverse step 11 took %s seconds" % (time.time() - start_time))
        
        # reverse step 9
        if self.debug:
            start_time = time.time()
        utils.dis(ag,undo=False)
        if self.debug:
            myprint("Reverse step 9 took %s seconds" % (time.time() - start_time))
        
        if self.use_pad_and_NN:
            # reverse step 8
            if self.debug:
                myprint('adjoint... Using padding and NN.')
                start_time = time.time()
            if self.use_float64:
                ag = torch.unsqueeze(torch.tensor(ag,dtype=torch.float64),dim=0)
            else:
                ag = torch.unsqueeze(torch.tensor(ag,dtype=torch.float32),dim=0)
            if self.cuda_avail:
                ag = ag.cuda()
            if self.debug:
                myprint("Reverse step 8 took %s seconds" % (time.time() - start_time))
            
            # reverse step 7
            if self.debug:
                start_time = time.time()
            ag = torch.autograd.grad(self.y, self.x, grad_outputs=ag, retain_graph=False)[0]
            if self.debug:
                myprint("Reverse step 7 took %s seconds" % (time.time() - start_time))
            
            # reverse step 6
            if self.debug:
                start_time = time.time()
            ag = torch.squeeze(ag).detach().cpu().numpy() 
            #ag = torch.squeeze(ag).detach().numpy()
            if self.debug:
                myprint("Reverse step 6 took %s seconds" % (time.time() - start_time))
        
            # reverse step 5
            if self.debug:
                start_time = time.time()
            
            ag = np.asarray(self.ag_pad(ag))[0] #not sure why adjoint outputs shape (1,3,128,128,128)
            if self.debug:
                myprint("Reverse step 5 took %s seconds" % (time.time() - start_time))
            
        else:
            myprint('adjoint... Skipping padding and NN.')

        # reverse step 4
        if self.debug:
            start_time = time.time()
        utils.dis(ag,undo=True)
        if self.debug:
            myprint("Reverse step 4 took %s seconds" % (time.time() - start_time))
        
        # reverse step 3
        if self.debug:
            start_time = time.time()
        self.ag_pos = ag.reshape(3,self.box.N[0]**3,order='C').T
        if self.debug:
            myprint("Reverse step 3 took %s seconds" % (time.time() - start_time))
        
        # memory
        if self.debug:
            start_time = time.time()
        self.ag_pos = np.copy(self.ag_pos,order='C')
        if self.debug:
            myprint("Fixing memory issue took %s seconds" % (time.time() - start_time))
            myprint('Finished backward pass of emu in %s seconds' % (time.time() - init_time))
        
    def getAdjointModel_impl(self, output_ag):
        # Set adjoint gradient wrt density field to zero
        output_ag[:] = 0
        # and instead specify the adjoint gradient wrt particle positions
        self.prev_module.adjointModelParticles(ag_pos=np.array(self.ag_pos, dtype=self.dtype), ag_vel=np.zeros_like(self.ag_pos, dtype=self.dtype))    
        
    @partial(jax.jit, static_argnums=(0,)) 
    def padding(self,x):
        return jnp.pad(x,((0,0),(48,48),(48,48),(48,48)),'wrap')
    
    @partial(jax.jit, static_argnums=(0,)) 
    def jax_cic(self, x, y, z):
        # TODO: fix hard-coded N and L
        Nx = Ny = Nz = 128
        Lx = Ly = Lz = 250.
        
        Ntot = Nx * Ny * Nz
        x = x * Nx / Lx
        y = y * Ny / Ly
        z = z * Nz / Lz

        qx, ix = utils.get_cell_coord(x)
        qy, iy = utils.get_cell_coord(y)
        qz, iz = utils.get_cell_coord(z)

        ix = ix.astype(int)
        iy = iy.astype(int)
        iz = iz.astype(int)
        rx = 1.0 - qx
        ry = 1.0 - qy
        rz = 1.0 - qz
        jx = (ix + 1) % Nx
        jy = (iy + 1) % Ny
        jz = (iz + 1) % Nz

        rho = jnp.zeros((Ntot,))

        for a in [False, True]:
            for b in [False, True]:
                for c in [False, True]:
                    ax = jx if a else ix
                    ay = jy if b else iy
                    az = jz if c else iz
                    ux = qx if a else rx
                    uy = qy if b else ry
                    uz = qz if c else rz

                    idx = az + Nz * ay + Nz * Ny * ax
                    rho += jnp.bincount(idx, weights=ux * uy * uz, length=Ntot)

        return rho.reshape((Nx, Ny, Nz)) / (x.shape[0] / Ntot) - 1.0
    
    @partial(jax.jit, static_argnums=(0,5,6,7,8,9,10)) 
    def cic_analytical_grad(self,density, x, y, z, Nx, Ny, Nz, Lx, Ly, Lz):
        Ntot = Nx * Ny * Nz
        x = x * Nx / Lx
        y = y * Ny / Ly
        z = z * Nz / Lz

        qx, ix = utils.get_cell_coord(x)
        qy, iy = utils.get_cell_coord(y)
        qz, iz = utils.get_cell_coord(z)

        ix = ix.astype(int)
        iy = iy.astype(int)
        iz = iz.astype(int)

        rx = 1.0 - qx
        ry = 1.0 - qy
        rz = 1.0 - qz
        jx = (ix + 1) % Nx
        jy = (iy + 1) % Ny
        jz = (iz + 1) % Nz

        adj_gradient = jnp.zeros((Nx**3,3))

        for coord in np.arange(3):
            if coord==0:
                rx = 1
                qx = -1
                ry = y - iy
                qy = 1 - ry
                rz = z - iz
                qz = 1 - rz

                adj_gradient = adj_gradient.at[:,0].set(density[ix,iy,iz] * qx * qy * qz + \
                                    density[ix,iy,jz] * qx * qy * rz + \
                                    density[ix,jy,iz] * qx * ry * qz + \
                                    density[ix,jy,jz] * qx * ry * rz + \
                                    density[jx,iy,iz] * rx * qy * qz + \
                                    density[jx,iy,jz] * rx * qy * rz + \
                                    density[jx,jy,iz] * rx * ry * qz + \
                                    density[jx,jy,jz] * rx * ry * rz)

            elif coord==1:
                rx = x - ix;
                qx = 1 - rx;
                ry = 1;
                qy = -1;
                rz = z - iz;
                qz = 1 - rz;

                adj_gradient = adj_gradient.at[:,1].set(density[ix,iy,iz] * qx * qy * qz + \
                                    density[ix,iy,jz] * qx * qy * rz + \
                                    density[ix,jy,iz] * qx * ry * qz + \
                                    density[ix,jy,jz] * qx * ry * rz + \
                                    density[jx,iy,iz] * rx * qy * qz + \
                                    density[jx,iy,jz] * rx * qy * rz + \
                                    density[jx,jy,iz] * rx * ry * qz + \
                                    density[jx,jy,jz] * rx * ry * rz)
            else:
                rx = x - ix;
                qx = 1 - rx;
                ry = y - iy;
                qy = 1 - ry;
                rz = 1;
                qz = -1;

                adj_gradient = adj_gradient.at[:,2].set(density[ix,iy,iz] * qx * qy * qz + 
                                    density[ix,iy,jz] * qx * qy * rz + 
                                    density[ix,jy,iz] * qx * ry * qz +
                                    density[ix,jy,jz] * qx * ry * rz +
                                    density[jx,iy,iz] * rx * qy * qz +
                                    density[jx,iy,jz] * rx * qy * rz +
                                    density[jx,jy,iz] * rx * ry * qz +
                                    density[jx,jy,jz] * rx * ry * rz)


        adj_gradient*= Nx / Lx

        return adj_gradient
    
    
class NullForward(borg.forward.BaseForwardModel):
    """
    BORG forward model which does nothing but stores
    the values of parameters to be used by the likelihood
    """
    def __init__(self, box: borg.forward.BoxModel) -> None:
        """
        Initialise the NullForward class
        Args:
            box (borg.forward.BoxModel): The input box model.
        """
        super().__init__(box, box)
        self.setName("nullforward")

        self.params = {}
        self.setCosmoParams(borg.cosmo.CosmologicalParameters())
        cosmo = self.getCosmoParams()
        cosmo.n_s = 0.96241
        self.setCosmoParams(cosmo)


    def setModelParams(self, params: dict) -> None:
        """
        Change the values of the model parameters to those given by params

        Args:
            params (dict): Dictionary of updated model parameters.
        """
        for k, v in params.items():
            self.params[k] = v
        print(" ")
        myprint(f'Updated model parameters: {self.params}')

    def getModelParam(self, model, keyname: str):
        """
        This queries the current state of the parameters keyname in model model.

        Args:
            model: The model
            keyname (str): The name of the parameter of interest
        """
        return self.params[keyname]