import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
import jax
import jax.numpy as jnp
import flax.linen as nn
from typing import Union
import sys
import os
import orbax
from orbax.checkpoint import PyTreeCheckpointer
import pickle
import numpyro.infer
import pandas as pd
import numpy as np
from functools import partial
from jax import jit
import ast
import arviz as az
from matplotlib import pyplot as plt
import math
import seaborn as sns
import argparse
import sys

print("Arguments received:", sys.argv)

argparser = argparse.ArgumentParser()
argparser.add_argument('--sites', type=str, nargs='+', help="list of site names (qvidja ruukki viikki)")
argparser.add_argument('--years', type=str, nargs='+', help="list of years (0,1,2,3)")
args = argparser.parse_args()

sites = args.sites
sites_years = {}
for site, year_site in zip(sites, args.years):
    years = [int(year) for year in year_site.split(',') if year.strip()]
    sites_years[site] = years
    print(sites_years)


print(years)
def load_checkpoint(ckpt_dir, scaler_X_file, scaler_y_file):
    ckpt = PyTreeCheckpointer().restore(ckpt_dir)
    state = ckpt['state']['params']
    with open(scaler_X_file,'rb') as f:
      scaler_X = pickle.load(f)
    with open(scaler_y_file,'rb') as f:
      scaler_y = pickle.load(f)
    
    return state, scaler_X, scaler_y

scaler_X_file = 'tmp/scaler_X_1.pkl'
scaler_y_file = 'tmp/scaler_y_1.pkl'
ckpt_dir = 'tmp/state_1'

model_params, scaler_X, scaler_y = load_checkpoint(ckpt_dir, scaler_X_file, scaler_y_file)


def sklearn_test( data, scaler_X, scaled): 
    b_df=pd.DataFrame(data)
    sklearn_scaled = scaler_X.transform(b_df)
    are_equal = jnp.allclose(scaled, sklearn_scaled, atol=0.001)

    print(are_equal) 

def jax_standard_scaler(X, scaler, task):
    mu = scaler.mean_  # numpy array of means
    sigma = scaler.scale_  # numpy array of standard deviations

    # convert mu and sigma to jax arrays
    mu = jnp.array(mu)
    sigma = jnp.array(sigma)

    if task == 'transform':
        return (X - mu) / sigma
    elif task == 'inverse':
        return (X * sigma) + mu
    else:
        raise ValueError("No given task.")
    
def reorder_features(met):
     original_columns = ['air_temperature', 'relative_humidity',
                      'surface_downwelling_shortwave_flux_in_air', 'precipitation_flux',
                      'wind_speed', 'harvest']
    
     new_columns = ['air_temperature', 'harvest',
                     'precipitation_flux', 'relative_humidity',
                      'surface_downwelling_shortwave_flux_in_air', 'wind_speed']

     # mapping from og order to alphabetical order
     order_map = {col: idx for idx, col in enumerate(new_columns)}


     reordered_met = jnp.zeros_like(met)
    
     for i, original_col in enumerate(original_columns):
         new_index = order_map[original_col]
         reordered_met = reordered_met.at[:, :, new_index].set(met[:, :, i])
 
     return reordered_met
     
def get_input(meteo, parameters, scaler_X): 
   
   both = jnp.concatenate([meteo, parameters], axis=1)
   print("here", both)

   scaled = jax_standard_scaler(both, scaler_X, 'transform')
   #sklearn_test( both, scaler_X, scaled)

   num_columns_meteo = meteo.shape[1]
   num_columns_parameters = parameters.shape[1]

   scaled_meteo = scaled[:, :num_columns_meteo]
   scaled_parameters = scaled[:, num_columns_meteo:num_columns_meteo + num_columns_parameters]

   #reshape meteo
   reshaped_met=scaled_meteo.reshape(num_points, 6, 53)

   reshaped_met=reshaped_met.transpose(0, 2, 1)

   reordered_met=reorder_features(reshaped_met)


   return reordered_met, scaled_parameters


class LSTM_model_simple(nn.Module):
   hidden_size_p1: Union[int, None]
   hidden_size_p2: Union[int, None]
   hidden_size_lstm: int
   hidden_size_comb: Union[int, None]

   @nn.compact
   def __call__(self , met, param):
      
      cell= nn.LSTMCell(name="lstmcell", features=self.hidden_size_lstm)

      def body_fn(cell, carry, met):
         carry, y= cell(carry, met)
         carry = (carry[0].astype(jnp.float32), carry[1].astype(jnp.float32))
         
         return carry, y
      
      # scan through each time step
      scan=nn.scan(body_fn, variable_broadcast="params",
            split_rngs={"params": False}, in_axes=1, out_axes=1)

      
      # concatenate parameters and meteorological inputs
      y_param = jnp.tile(param[:, jnp.newaxis, :], (1, 53, 1))

      y = jnp.concatenate([met, y_param], axis=-1)


      # FNN layers
      if self.hidden_size_p1 is not None:
          y=nn.Dense(self.hidden_size_p1) (y)
          y=nn.relu(y)

      if self.hidden_size_p2 is not None:
          y=nn.Dense(self.hidden_size_p2)(y)
          y=nn.relu(y)

      # LSTM layer
      carry=cell.initialize_carry(jax.random.PRNGKey(0),  ((y.shape)[0], self.hidden_size_lstm))
      carry, y = scan(cell, carry, y)
      y=nn.relu(y)

      # FNN layer
      if self.hidden_size_comb is not None:
          y=nn.Dense(int(self.hidden_size_comb))(y)
          y=nn.relu(y)
      
      # FNN layer to produce final output
      y=nn.Dense(5)(y)
      
      return y

# give optimized hyperparameters
model = LSTM_model_simple(hidden_size_lstm=162, hidden_size_p1=None, hidden_size_p2=None, hidden_size_comb=211)
print(model)


def calibration_model(model_params, meteo, parameters,y=None):


    parameter_priors = {
        'LAICR': dist.Uniform(2.0, 7.0),
        'LAIEFT': dist.Uniform(0.1, 0.35),
        'LAITIL': dist.Uniform(0.3, 1.1),
        'NELLVM': dist.Uniform(1.0, 2.75),
        'PHY': dist.LogNormal(4.25, 0.25),
        'ROOTDM': dist.Uniform(0.5, 1.5),
        'RUBISC': dist.LogNormal(loc=0.95, scale=0.66),
        'SHAPE': dist.Beta(4.0, 2.0),
        'TBASE': dist.Uniform(3.0, 10.0),
        'TCRES': dist.Uniform(1.0, 3.0)
    }

    # sample each parameter
    for i, (param_name, prior) in enumerate(parameter_priors.items()):
        sampled_value = numpyro.sample(param_name, prior)
        parameters = parameters.at[:, i].set(sampled_value)
        print(f"{param_name} sampled value: {sampled_value}")

    sigma = numpyro.sample('sigma', dist.InverseGamma(1, 1))
    
    def predict(model_params, meteo, parameters):
        
        met, param=get_input(meteo, parameters, scaler_X)
        print(param)

        # predict with emulator and scale back to original scale
        predicted=model.apply({'params': model_params}, met, param)
        pred_shape=predicted.shape
        pred_scale_shape = predicted.reshape((predicted.shape[0], -1))

        # inverse transform the scaled predictions to original scale
        predicted=jax_standard_scaler(pred_scale_shape, scaler_y, 'inverse')
        predicted=predicted.reshape(pred_shape)
        
        return predicted
    
    # generate predictions and filter only gpp values
    preds = predict(model_params, meteo, parameters)
    preds_gpp=preds[:,:,1]
    
    # flatten predictions and observations for likelihood computation
    preds_gpp_flat = preds_gpp.flatten()
    y_flat = y.flatten()
    
    valid_mask = ~jnp.isnan(y_flat)

    with numpyro.plate('data', preds_gpp_flat.shape[0]):
        with numpyro.handlers.mask(mask=valid_mask):
            numpyro.sample('obs', dist.Normal(preds_gpp_flat, sigma), obs=y_flat)


# get GPP data for calibration
if  sites=='all':
    sites = ['qvidja', 'ruukki', 'viikki']

all_data = []

for s in args.sites:
    years=sites_years[s]
    print(f"Processing site: {s}, years {years} ")
    # import data from given site
    data = pd.read_parquet(f'GPP_data/joined_df_{s}.parquet').iloc[years]
    print(data)
    data['site'] = s
    print(f"Selected data for {s}:")
    print(data)


    all_data.append(data)
    
data = pd.concat(all_data, ignore_index=True)

num_points=data.shape[0]

# GPP data
gpp_df=data.GPP
gpp_array=jnp.array(gpp_df)

y = gpp_array.reshape((num_points, 53, 1))

print("Y:", y)

# meteorological inputs
meteo_df = data.meteo

soilparams_df = data.soilparam

# initalize the parameter dic
other_params = {
    'LAICR':    np.nan,
    'LAIEFT':   np.nan,
    'LAITIL':   np.nan,
    'NELLVM':   np.nan, 
    'PHY':      np.nan,
    'ROOTDM':   np.nan,
    'RUBISC':   np.nan,
    'SHAPE':    np.nan,
    'TBASE':    np.nan,
    'TCRES':    np.nan 
    }

# data processing
param_df= pd.DataFrame([other_params] , columns=other_params.keys())
param_df = pd.concat([param_df]*num_points, ignore_index=True)

df_all_param = pd.concat([param_df, soilparams_df], axis=1, ignore_index=True)

meteo=jnp.array(meteo_df)
parameters=jnp.array(df_all_param)


# initialize the NUTS HMC
nuts_kernel = numpyro.infer.NUTS(calibration_model)

mcmc = numpyro.infer.MCMC(nuts_kernel, num_warmup=10, num_samples=100, num_chains=1) 

mcmc.run(jax.random.PRNGKey(0),model_params, meteo, parameters, y)

samples = mcmc.get_samples() #retrieves the posterior distribution

numpyro.diagnostics.print_summary(
    samples,
    prob=0.9,
    group_by_chain=False
)

# convert samples to ArviZ InferenceData
idata = az.from_numpyro(mcmc)

# traceplot
az.plot_trace(idata)

file_name_sites = "_".join(sites)
plt.savefig(f'{file_name_sites}/trace_plot_{years}.png') 
plt.show()

# load samples to a pickle file
with open(f'MCMC_samples/{file_name_sites}/samples_{years}.pkl', 'wb') as f:
    pickle.dump(samples, f)