""" 
Causal Confusion CEM Planner.
Test if causal curiosity can help distinguish between two causal factors
not introduced in training, but in testing.
"""
import random
import time
import os
from datetime import datetime
from pathlib import Path
from functools import partial
import gym
from os.path import exists

import torch 
from sklearn import preprocessing
from numpy import linalg as LA

import numpy as np
import matplotlib.pyplot as plt
#from sklearn.cluster import KMeans
from tslearn.clustering import TimeSeriesKMeans, silhouette_score
from loguru import logger
from tqdm import tqdm
import dill

from gym.wrappers.monitoring.video_recorder import VideoRecorder

import multiprocessing_on_dill as mp
import psutil
# local imports
from cem_planner import CEMPlanner
from plan_action_spaces import get_plan_action_space
from cem.uniform_bounds import UniformBounds
from causal_world.task_generators.task import task_generator
from causal_world.envs.causalworld import CausalWorld


seeds = [122, 123, 124, 125, 126, 127, 233, 344, 455, 566, 677, 788, 889, 900, 1111]
masses = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
sizes = [0.05, 0.055, 0.06, 0.065, 0.07, 0.075, 0.08, 0.085, 0.095] # median 0.75, count 9
frame_skips = [1,2,3,4,5,6,7,8,9]
frictions = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
dampings = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
action_mode = 'RGB_asym'
obs_dim = 3


""" Implements Probabilistic Neural Networks using PyTorch
    The output of a PNN defines a mean and a variance of a Gaussian.
    Then, a sample is generated by sampling from the distribution
    defined by this mean and variance.

    See:
    https://arxiv.org/abs/1612.01474
"""

class Model():
    """ Contains PyTorch model """
    def __init__(self, seed):
        torch.manual_seed(seed)
        self.input_dim = 12
        self.hidden_units = 16
        self.output_dim = 6

        device = torch.device('cpu')
        # Define the model [1 x h x h x 2]
        self.model = torch.nn.Sequential(
            torch.nn.Linear(self.input_dim, self.hidden_units, bias=True),
            torch.nn.ReLU(),
            torch.nn.Linear(self.hidden_units, self.hidden_units, bias=True),
            torch.nn.ReLU(),
            torch.nn.Linear(self.hidden_units, self.hidden_units, bias=True),
            torch.nn.ReLU(),
            torch.nn.Linear(self.hidden_units, self.output_dim, bias=True),
        ).to(device)
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-3)
        # self.optimizer = torch.optim.SGD(self.model.parameters(), lr=1e-2, momentum=0.5)
        self.losses = []
        self.means = []
        self.var = []
        self.max_var = 0
        self.min_var = float('inf')

    def softplus(self, x):
        """ Positivity constraint """
        softplus = torch.log(1+torch.exp(x))
        # Avoid infinities due to taking the exponent
        softplus = torch.where(softplus==float('inf'), x, softplus)
        return softplus

    def adjust_learning_rate(self):
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = param_group['lr']*0.999

    def NLL(self, means, var, truth):
        """ Compute the Negative Log Likelihood """
        diff = torch.sub(truth, means)
        var = self.softplus(var)
        # Check min and max variance
        for v in var:
            if v == float('inf'):
                raise ValueError('Infinite variance')
            if v > self.max_var:
                self.max_var = v 
            if v < self.min_var:
                self.min_var = v
        loss = torch.mean(torch.div(diff**2, var))
        loss += torch.mean(torch.log(var))
        return loss 

    def forward(self, inputs):
        """ Forward pass for given inputs """
        x = torch.from_numpy(inputs)
        # Compute output of the trained model
        model_out = self.model(x)
        mean, var = torch.split(model_out, 3, dim=0)
        # Bound the variance
        var = self.softplus(var)
        # TRICK -- not added until now, actually increases overall variance
        # Avoids exploding std
        # logvar = torch.log(var)
        # logvar = self.max_var - self.softplus(self.max_var-logvar)
        # logvar = self.min_var + self.softplus(logvar-self.min_var)
        # var = torch.exp(logvar)
        return mean.detach().numpy(), var.detach().float().numpy()

    def step(self, inputs, labels):
        """ Execute a single gradient step using the given samples """
        # Convert samples to useable tensors
        # - need to unsqueeze since PyTorch only takes mini-batches
        #   we need to add an additional dimension
        # - for some reason we also need to call .float()
        #   even though x and y are already of type torch.float64
        x = torch.from_numpy(inputs)
        y = torch.from_numpy(labels)

        # Compute the output of the network
        model_out = self.model(x)
        means, var = torch.split(model_out, 3, dim=0)
        var = self.softplus(var)
        self.means.append(means)
        self.var.append(var)
#         print(means)
#         print(var)
        # Compute the NNL loss
        nll = self.NLL(means, var, y)
        self.losses.append(nll)

        # Do the backward pass through the network
        self.optimizer.zero_grad()
        nll.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5)
        self.optimizer.step()

        self.adjust_learning_rate()

        
        
def KL(meana, meanb, cova, covb):
    return (1/2)*(np.log(LA.norm(covb)/LA.norm(cova)) - len(meana) + np.matmul(np.matmul(np.transpose(meanb - meana), LA.inv(covb)), (meanb - meana)) + np.trace(np.matmul(LA.inv(covb), cova)))

def pnnOut(X, y):
    
    model = Model(12)
    for ep in range(10):
        for i in range(len(X)):
            model.step(X[i],y[i])
    means, var = model.forward(X[0])
    varmat = np.diag(var)

       
    return means, varmat
   
def processXy(states, actions):
    X = []
    y = []
    for i in range(len(states)):

        s = np.array(states[i][:-1])
        label = np.array(states[i][1:])
        a = np.array(actions[i][1:])
        feature = np.concatenate([s,a],axis = 1)
        X.append(feature)
        y.append(label)
    X = np.array(X)
    y = np.array(y)
    X = X.reshape((-1, 12))
    y = y.reshape((-1, 3))

    X = np.float32(X)
    y = np.float32(y)
    X = preprocessing.minmax_scale(X)
    y = preprocessing.minmax_scale(y)


    return X, y    

class CEMPlannerCW(CEMPlanner):
    """
    Extension of CEMPlanner class.
    New plan function that allows us to specify which env to train in.
    Returns: best-trained action plan, observations, trained KMeans model
    """

    # plan envs 
    def getEnvs(self, masses = [0.5], sizes = [0.07], frictions = [0.1], dampings = [0.5], frame_skips = [1], bsize = 0.07, bmass = None, bfriction = 0.1, bdamping = 0.5, bframe_skip = 1
        ):
        
        envs = []

        for mass in masses:
            for size in sizes:
                for friction in frictions:
                    for damping in dampings:
                        for frame_skip in frame_skips:
                            task = task_generator(task_generator_id ='lifting',
                                tool_block_mass = mass,
                                tool_block_shape = 'cube',
                                tool_block_size = size)
                            env = CausalWorld(task=task, 
                                skip_frame = self.frame_skip,
                                enable_visualization = False)
                            env.setFrictionDamping(f =friction, d= damping)
                            envs.append(env)
        if bmass != None:
            task = task_generator(task_generator_id ='lifting',
                        tool_block_mass = bmass,
                        tool_block_shape = 'cube',
                        tool_block_size = bsize)

            env = CausalWorld(task=task, 
                        skip_frame = bframe_skip, 
                        enable_visualization = False)

            env.setFrictionDamping(f = bfriction, d=bdamping)
            envs.append(env)

        return envs

    # plan envs with test env
   

    # plan envs with test env
   


# for multiprocessing
def worker(i):
    p = psutil.Process()
    # Params. Play around with these settings. They are not yet optimized.
    #scenario = 'lift'
    p.cpu_affinity([i])
    matrix = np.zeros((9,9))


    total_budget = 400
    plan_horizon = 6
    n_plan_iterations = 20
    sampler = 'uniform'
    warm_starts = False
    warm_start_relaxation = 0.0
    elite_fraction = 0.1
    viz_progress = True
    n_frames = 198
    seed = seeds[i]
    rng = np.random.RandomState(seed)
    np.random.seed(seed)

    plan_action_repeat = np.floor_divide(n_frames, plan_horizon)
    n_plan_cache_k = plan_horizon
    n_plans = np.floor_divide(total_budget * n_plan_cache_k,
                              plan_horizon * n_plan_iterations)
    n_plan_elite = max(1, int(round(elite_fraction * n_plans)))

    if n_plans <= n_plan_elite:
        n_plan_elite = n_plans - 1

    # action_space is gym `Box` env that defines vals for each action [-1,1]
    # action_transf is a func that returns an array for real-val actions? 
    action_space, action_transformation = get_plan_action_space(action_mode) 


    # load pre-calculated rewards and prediction
   
    rng = np.random.RandomState(seed)
    np.random.seed(seed)
    print(f"start training planner {i+1}")
    training_planner = CEMPlannerCW(n_plans=n_plans,
                           horizon=plan_horizon,
                           action_space=action_space,
                           sampler=sampler,
                           n_iterations=20,
                           n_elite=n_plan_elite,
                           cache_k=n_plan_cache_k,
                           obs_dim = obs_dim,
                           warm_starts=warm_starts,
                           warm_start_relaxation=warm_start_relaxation,
                           plan_action_repeat=plan_action_repeat,
                           action_transformation=action_transformation,
                           rng=rng,
                           viz_progress=viz_progress,
                           )

    envs = training_planner.getEnvs(sizes = sizes)
    print("num of training envs: ", len(envs))

    training_action_plan, training_rel_duration_plan, training_observations, training_km_sdtw, training_best_return = training_planner.plan(envs)
    

    trainingstates = []
    trainingactions = []
    for i_env, env in enumerate(envs):
        state, action = training_planner.simulate_and_write(env, training_action_plan, training_rel_duration_plan)
        trainingstates.append(state)
        trainingactions.append(action)
    
    X_training, y_training = processXy(trainingstates, trainingactions)
    mean_training, cov_training = pnnOut(X_training, y_training)  

   
    for i_size, size in enumerate(sizes):
        for i_mass, mass in enumerate(masses):
            envs = training_planner.getEnvs(sizes = [size], masses = [mass])          
            states, actions = training_planner.simulate_and_write(envs[0], training_action_plan, training_rel_duration_plan)
            X, y = processXy([states], [actions])
            mean, cov = pnnOut(X, y)  
            kl = KL(mean, mean_training, cov, cov_training)
            
            matrix[i_size][i_mass] = kl 
            print(f"{i_size}_{i_mass}_{i}")
    print(matrix)
    dill.dump(matrix, open(f'MassFrameskip/SA/kl_{i}.pickle', 'wb'))


def main(output_dir):

    # change into output directory
    if not output_dir.exists():
        os.mkdir(output_dir)
    os.chdir(output_dir)
    # Params. Play around with these settings. They are not yet optimized.
    #scenario = 'lift'
    manager = mp.Manager()
    print("cpu count: ", mp.cpu_count())
    pool = mp.Pool(16)

    jobs = []
    for i in range(15%mp.cpu_count()):
        job = pool.apply_async(worker, (i, ))
        jobs.append(job)
        time.sleep(1)

    for job in jobs:
        job.get()

    pool.close()
    pool.join()


  


if __name__ == '__main__':
    print("CEM Planner CausalWorld")

    ## change this line to store your data in whatever folder you'd like
    output_dir = Path('./pickle/CW')
    print(f"Output dir: {output_dir}")
    print("Make sure this directory is what you expected, if not change it!")
    main(output_dir=output_dir)



    



