#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jul  5 00:27:15 2023

@author: chen
"""

from PyRAI2MD.variables import ReadInput

from photoMD.settings import generator_setting
from photoMD.tools import md_tools

import numpy as np
try:
    from importlib.resources import files
except:
    from importlib_resources import files
import os, threading, shutil

HtoEv = 27.21138624598853
AToBohr = 1.889725989

class UserGene(object):
    """
    User defined Generator. Receive prediction from Prediction kernel and generate new data points.
    """
    def __init__(self, rank, result_dir):
        """
        initilize the generator.
        
        Args:
            rank (int): current process rank (PID).
            result_dir (str): path to directory to save metadata and results.
        """
        self.rank = rank
        self.result_dir = os.path.abspath(result_dir)
        ##### User Part ######
        self.job_keywords = ReadInput(generator_setting.variable_input)    # input for MD trajectory setting
        self.traj_count = 0    # count the trajectories generated by each generator process
        self.n_state = int(self.job_keywords["molecule"]["ci"][0])    # number of electronic states for simulation
        self.n_atom = len(generator_setting.elements)    # number of atoms of the target molecule
        self.save_path = os.path.join(self.result_dir, f"traj_data_rank_{self.rank}.pickle")

        # load initial conditions
        with open(files("photoMD.data").joinpath("initial_conditions_small_set.npy"), 'rb') as fh:
            self.init_coord = np.load(fh)
            self.init_velc = np.load(fh)
            self.root = np.load(fh)
        self.root += 1
        self.init_idx = np.arange(0, self.init_coord.shape[0], dtype=int)

        # TODO: restart from breaking point

        # record trajectory data
        self.step_keys = ['coord', 'energy', 'force', 'state']    # list of keys of traj_data to data related to each time step
        self.traj_keys = ['termin',]    # list of keys of traj_data to data related to trajectories
        self.traj_data = {
            'energy': [],
            'force': [],
            'coord': [],
            'state': [],
            'termin': [],
            }
        
        os.chdir(self.result_dir)
        
    def generate_new_data(self, data_to_gene):
        """
        Generate new data point based on data_to_gene (prediction from Prediction kernel).
        
        Args:
            data_to_gene (1-D numpy.ndarray or None): data from prediction kernel through EXCHANGE process.
                                                      Initialized as None for the first time step.
                                                      Source: element of data_to_gene_list from UserModel.predict()
            
        Returns:
            stop_run (bool): flag to stop the active learning workflow. True for stop.
            data_to_pred (1-D numpy.ndarray): data to prediction kernel through EXCHANGE process.
                                              Destination: element of input_list at UserModel.predict()
        """
        stop_run = False
        data_to_pred = None
        
        # please notice that data_to_gene is intinilized to be None for the first iteration.
        ##### User Part #####
        traj_status = 0
        # initialize new trajectory if data_to_gene is None
        if data_to_gene is None:
            self.aimd = self._init_traj()
            current_state, coord = self.aimd.propagate_step_one()

        # unpack received data to energy and force predictions
        # # shape of energy: (n_state,)
        # # shape of force: (n_state, n_atom, 3)
        else:
            eng_pred = data_to_gene[:self.n_state].reshape(self.n_state,) / HtoEv
            force_pred = data_to_gene[self.n_state:].reshape(self.n_state, self.n_atom, 3) / (HtoEv * AToBohr)

            # check eng_pred and force_pred
            # # if all zeros, then this is a stop signal due to the high STD of predictions
            # # stop the current trajectory, save the termination reason, and initialize a new trajectory
            if (eng_pred == 0).all() or (force_pred == 0).all():
                traj_status = 4    # indication of termination because of high STD

            # # if receive normal energy and force predictions, progate the trajectory to next step
            else:
                traj_status = self.aimd.propagate_step_two(eng_pred, force_pred)
                self.aimd.traj.iter += 1
                self.traj_data["coord"][-1].append(np.copy(self.aimd.traj.coord))
                self.traj_data["energy"][-1].append(np.copy(self.aimd.traj.energy))
                self.traj_data["force"][-1].append(np.copy(self.aimd.traj.grad))
                self.traj_data["state"][-1].append(self.aimd.traj.state - 1)
                current_state, coord = self.aimd.propagate_step_one()
        
        # check trajectory status and the number of steps
        # # if trajectory status is not 0, then this is a stop signal due to abnormal bond length
        # # if the number of steps exceeds the max steps, stop the current trajectory
        # # then start a new trajectory
        if traj_status != 0 or self.aimd.traj.iter > self.aimd.step_left:
            self.traj_data["termin"].append(traj_status)
            self.traj_count += 1
            os.chdir(self.result_dir)
            shutil.rmtree(self.traj_dir)
            self.aimd = self._init_traj()
            current_state, coord = self.aimd.propagate_step_one()

        # organize the generate coordinates and electronic state as a 1-D numpy array to send to the Prediction kernel
        data_to_pred = np.append([float(current_state),], coord.flatten(), axis=0)
        stop_run = False

        # check the number of trajectories generated by the generator process
        # # if exceeds the limit, stop the PAL workflow
        if self.traj_count >= generator_setting.max_traj_per_process:
            data_to_gene = np.zeros((1+self.n_atom*3,), dtype=float)
            stop_run = True

        # stop_run should be returned as a bool value
        # data_to_pred should be returned as an 1-D numpy array
        return stop_run, data_to_pred
    
    def save_progress(self, stop_run):
        """
        Save the current state and progress. Called everytime after the interval defined by progress_save_interval in al_setting, and when the active learning workflow is shutdown (stop_run is True).

        Args:
            stop_run (bool): flag to stop the active learning workflow. True for stop.
        """
        ##### User Part #####
        mode = 'ab' if os.path.exists(self.save_path) else 'wb'
        save_thread = threading.Thread(target=md_tools.save_pickle, name=f"save_traj_{self.rank}",\
                                    args=(self.traj_data, self.step_keys, self.traj_keys,\
                                          [None,self.save_path], mode, self.aimd.traj.coord,\
                                        self.aimd.traj.state, None, self.aimd.traj.grad, None), daemon=True)
        save_thread.start()

        self.traj_data = {
                        'energy': [[],],
                        'force': [[],],
                        'coord': [[],],
                        'state': [[],],
                        'termin': [],
                        }

    def stop_run(self):
        """
        Called before the Generator process terminating when active learning workflow shuts down.
        """
        ##### User Part #####
        self.save_progress(stop_run=True)

    def _init_traj(self):
        """
        Initialize new MD trajectory.

        Args:
            None

        Returns:
            aimd (AIMD object): Object managing a MD trajectory.
        """
        # create directory for each trajectory to store meta data
        self.traj_dir = f"traj_{self.rank}"
        os.makedirs(self.traj_dir, exist_ok=True)
        os.chdir(self.traj_dir)

        # setup trajectory data record
        for k in self.traj_data.keys():
            if k != 'termin':
                self.traj_data[k].append([])

        # randomly sample initial conditions
        i = np.random.choice(self.init_idx, size=1)
        while self.job_keywords['md']['root'] is None and int(self.root[i]) != int(self.job_keywords['md']['root']):
            i = np.random.choice(self.init_idx, size=1)
        self.job_keywords['md']['root'] = int(self.root[i])
        # set up MD trajectory
        mol = md_tools.read_inicond(self.init_coord[i], self.init_velc[i], generator_setting.elements)[0]
        aimd = md_tools.set_aimd(mol, self.job_keywords, generator_setting.bond_index, generator_setting.bond_limit)
        aimd.traj.iter = 1
        return aimd
