from ..model import Model
import json, os
from abacustest.lib_prepare.abacus import WriteKpt, WriteInput, ReadInput, AbacusStru
from abacustest.lib_prepare.comm import kspacing2kpt
from abacustest.constant import RECOMMAND_IMAGE, RECOMMAND_COMMAND, RECOMMAND_MACHINE

from typing import List, Dict, Any
from typing import Union
import warnings
import shutil, glob, copy, traceback
from abacustest.lib_collectdata.collectdata import RESULT

from abacustest.outresult import pandas_out


class ElasticModel(Model):
    @staticmethod
    def model_name(): # type: ignore
        '''
        Name of the model, which will be used as the subcommand
        '''
        return "elastic"
    
    @staticmethod
    def description(): # type: ignore
        '''
        Description of the model
        '''
        return "Prepare and postprocess the elastic"
    
    @staticmethod
    def prepare_args(parser):
        '''
        Add arguments for the prepare subcommand
        The arguments can not be command, model, modelcommand 
        '''
        parser.description = "Prepare the inputs for the elastic calculation."
        parser.add_argument('-j', '--job', default=[], action="extend", nargs="*", help='the paths of ABACUS jobs, should contain INPUT, STRU, or KPT, and pseudopotential and orbital files')
        parser.add_argument("--norm", default=0.01, type=float, help="The maximum strain for normal mode, default is 0.01")
        parser.add_argument("--shear", default=0.01, type=float, help="The maximum strain for shear mode, default is 0.01")
        parser.add_argument("--norelax", action="store_true", help="Whether to not do atomic relaxation for each deformed structure, default is doing.")
        parser.add_argument("--image", type=str, default=RECOMMAND_IMAGE, help="The image to use for the Bohrium job, default is %s" % RECOMMAND_IMAGE)
        parser.add_argument("--machine", type=str, default=RECOMMAND_MACHINE, help="The machine to use for the Bohrium job, default is 'c32_m64_cpu'.")
        parser.add_argument("--abacus_command", type=str, default=RECOMMAND_COMMAND, help=f"The command to run the Abacus job, default is '{RECOMMAND_COMMAND}'.")
        return parser
    
    
    def run_prepare(self,params):
        '''
        Run the model with the given parameters
        '''
        if not params.job:
            raise ValueError("No job specified, please use -j or --job to specify the job paths.")
        
        paths = PrepElastic(params.job, norm=params.norm, shear=params.shear, norelax=params.norelax).run()
        
        setting = {
            "save_path": "results",
            "run_dft": []
        }
        
        for path in paths:
            setting["run_dft"].append({
                "ifrun": True,
                "example": path,
                "command": params.abacus_command,
                "image": params.image,
                "bohrium": {
                    "scass_type": params.machine,
                    "job_type": "container",
                    "platform": "ali"
                }
            })
        
        setting_file = "setting.json"
        json.dump(setting, open(setting_file, "w"), indent=4)
        print("\nThe inputs are generated in", ", ".join(params.job))
        print(f"You can modify '{setting_file}', and execute below command to run the abacustest to submit all jobs to bohrium:\n\tabacustest submit -p setting.json\n")
        print(f"After finishing the calculations, you can enter the results directory, \nand run below command below to postprocess the elastic results:\n\tabacustest model elastic post -j {' '.join(params.job)}\n")
        
    
    def postprocess_args(parser):
        '''
        Add arguments for the postprocess subcommand
        The arguments can not be command, model, modelcommand 
        '''
        parser.description = "Postprocess the elastic calculation results."
        parser.add_argument('-j', '--job', default=[], action="extend", nargs="*", help='the paths of the job directories, should contain the results of deformed structures generated by the prepare step.')
        return parser   

    def run_postprocess(self,params):
        '''
        Parse the parameters and run the postprocess process'''
        if not params.job:
            raise ValueError("No job specified, please use -j or --job to specify the job paths.")
        
        metrics, results = PostprocessElastic(params.job).run()
        json.dump(metrics, open("metrics.json", "w"), indent=4)
        json.dump(results, open("metrics_elastic.json", "w"), indent=4)
        pandas_out(results, "metrics_elastic.csv",print_list_seperate=True)
        print("\nThe postprocess is done. The metrics are saved in 'metrics.json', and the elastic results are saved in 'metrics_elastic.json'.")
        


class PrepElastic():
    """Prepare the elastic calculation.
    
    Args:
        job (str or list of str): the path of the job directory or a list of job directories.
        norm (float): the maximum strain for normal mode, default is 0.01
        shear (float): the maximum strain for shear mode, default is 0.01
        relax (bool): whether to do atomic relaxation for each deformed structure, default is False
    
    run() function can generate the deformed structures and write it into the job directories, and also link the original files to the deformed directories.
    It will return a list of paths of the deformed directories.
    The deformed directories will be named as "deformed_00", "deformed_01", ..., "deformed_N" where N is the number of deformed structures.
    """
    def __init__(self, job: Union[str, List[str]], 
                 norm: float = 0.01, 
                 shear: float = 0.01,
                 norelax: bool = False):
        self.job = job if isinstance(job, list) else [job]
        self.norm_strain = norm
        self.shear_strain = shear
        self.norelax = norelax
    
    def run(self):
        paths = []
        for job in self.job:
            if not os.path.isdir(job):
                warnings.warn(f"Job {job} is not a directory, skip it.")
                continue
            print(f"Preparing elastic calculation for job: {job}")
            paths.append(self.prepare_one_job(job))
        
        return paths 
            
    
    def prepare_one_job(self, job: str):
        '''
        Prepare the elastic calculation for one job
        '''
        # clean the deformed paths
        shutil.rmtree(os.path.join(job,"org"), ignore_errors=True)
        for i in glob.glob(os.path.join(job, "deformed_*")):
            shutil.rmtree(i, ignore_errors=True)
        
        stru = AbacusStru.ReadStru(os.path.join(job, "STRU"))
        
        # modify INPUT file
        input_param = ReadInput(os.path.join(job, "INPUT"))
        if self.norelax:
            input_param["calculation"] = "scf"
        else:
            input_param["calculation"] = "relax"
        input_param["cal_stress"] = 1
            
        if "kspacing" in input_param:
            kpt = kspacing2kpt(float(input_param["kspacing"]), stru.get_cell(bohr=True))
            warnings.warn(f"Transform kspacing {input_param['kspacing']} to kpt {kpt}.")
            WriteKpt(kpt, os.path.join(job, "KPT"))
            del input_param["kspacing"]

        org_files = os.listdir(job)
        for ifile in ["INPUT", "STRU"]:
            if ifile in org_files:
                org_files.remove(ifile)   
        
        # generate the deformed structures
        deformed_strus, strains = self._gen_deformed_structure(stru)

        org_path = os.path.join(job, "org")
        self._gen_subjob(job, org_path, org_files, input_param, stru)
        
        job_paths = [org_path]
        deformed_paths = []
        for i, def_stru in enumerate(deformed_strus):
            deformed_path = f"deformed_{i:02d}"
            ipath = os.path.join(job, deformed_path)
            self._gen_subjob(job, ipath, org_files, input_param, def_stru)
            job_paths.append(ipath)
            deformed_paths.append(deformed_path)
        
        json.dump({"strains": [s.as_dict() for s in strains],
                   "deformed_paths": deformed_paths}, 
                    open(os.path.join(org_path, "elastic.json"), "w"),
                          indent=4)

        return job_paths    

    def _gen_deformed_structure(self, stru: AbacusStru):
        from pymatgen.analysis.elasticity.strain import DeformedStructureSet
        from pymatgen.analysis.elasticity.elastic import Strain
        
        structure = stru.to_pymatgen()
        
        norm_strains = [-self.norm_strain, -0.5*self.norm_strain, +0.5*self.norm_strain, +self.norm_strain]
        shear_strains = [-self.shear_strain, -0.5*self.shear_strain, +0.5*self.shear_strain, +self.shear_strain]
        
        
        deformed_strus = DeformedStructureSet(structure,
                                    symmetry=False,
                                    norm_strains=norm_strains,
                                    shear_strains=shear_strains)
        
        strus = []
        for defor in deformed_strus:
            new_stru = copy.deepcopy(stru)
            new_stru.set_cell(defor.lattice.matrix, bohr=False)
            new_stru.set_coord(defor.frac_coords, bohr=False, direct=True)
            strus.append(new_stru)
        
        strains = [Strain.from_deformation(d) for d in deformed_strus.deformations]  
        
        return strus, strains
    
    def _gen_subjob(self, 
                    job: str,
                    ipath: str, 
                    org_files: List[str],
                             input_param: Dict[str, Any],
                             new_stru: AbacusStru):
        """Generate a subjob for the elastic calculation.
        
        Args:
            job (str): the path of the job directory.
            ipath (str): the path of the subjob directory.
            org_files (List[str]): the list of original files to link.
            input_param (Dict[str, Any]): the input parameters for the subjob.
            new_stru (AbacusStru): the new structure for the subjob.
        """
        os.makedirs(ipath, exist_ok=True)
        new_stru.write(os.path.join(ipath, "STRU"))
        WriteInput(input_param, os.path.join(ipath, "INPUT"))
        
        # link the original files
        for ifile in org_files:
            os.symlink(os.path.abspath(os.path.join(job, ifile)), os.path.join(ipath, ifile))
        
        

class PostprocessElastic():
    """Postprocess the elastic calculation results.
    
    Args:
        job (str or list of str): the path of the job directory or a list of job directories. 
            Each job should contain directories "org" and "deformed_00", "deformed_01", ..., "deformed_N" where N is the number of deformed structures.
            "org" is the results of the original structure, and a file elastic.json should be present in the "org" directory, which contains the strain values.

    """
    def __init__(self, job: Union[str, List[str]]):
        self.job = job if isinstance(job, list) else [job]
    
    def run(self):
        results = {}
        metrics = {}
        for job in self.job:
            if not os.path.isdir(job):
                warnings.warn(f"Job {job} is not a directory, skip it.")
                continue
            print(f"Postprocessing elastic calculation for job: {job}")
            try:
                m, r = self.postprocess_one_job(job)
            except Exception as e:
                traceback.print_exc()
                warnings.warn(f"Failed to postprocess job {job}: {e}")
                continue
            
            metrics.update(m)
            results[job] = r
        
        return metrics, results
    
    def postprocess_one_job(self, job: str):
        '''
        Postprocess the elastic calculation for one job
        '''
        # check files
        elastic_file = os.path.join(job, "org", "elastic.json")
        original_stru_file = os.path.join(job, "org", "STRU")
        if not os.path.isfile(elastic_file) or not os.path.isfile(original_stru_file):
            raise ValueError(f"Elastic file {elastic_file} or original structure file {original_stru_file} not found in {job}.")
        
        sub_dirs = [os.path.join(job, "org")] + [os.path.join(job, f"deformed_{i:02d}") for i in range(24)]
        
        metrics = self.get_metrics(sub_dirs)
        strain_param = json.load(open(elastic_file))
        strains = strain_param["strains"]
        deformed_paths = strain_param["deformed_paths"]
        
        input_stress = self.abacus_stress_to_pymatgen_stress(metrics[os.path.join(job, "org")]["stress"])
        deformed_stress = [self.abacus_stress_to_pymatgen_stress(metrics[os.path.join(job, dp)]["stress"]) for dp in deformed_paths]
        elastic_results = self.cal_elastic(strains, deformed_stress, input_stress)
        
        return metrics, elastic_results
        
    def abacus_stress_to_pymatgen_stress(self, stress: List[float]):
        """
        Transform calculated stress (units in kBar) collected by abacustest
        to Pymatgen format (units in GPa)
        """
        from pymatgen.analysis.elasticity.stress import Stress
        import numpy as np
        
        return Stress(-0.1 * np.array([stress[0:3],
                                   stress[3:6],
                                   stress[6: ]])) # 1 kBar = 0.1 GPa    
        
    def cal_elastic(self, strain, deformed_stru_stresses, input_stru_stress):
        from pymatgen.analysis.elasticity.elastic import ElasticTensor
        from pymatgen.analysis.elasticity.elastic import Strain
        strain = [Strain.from_dict(s) for s in strain]
        
        result = ElasticTensor.from_independent_strains(strain,
                                                    deformed_stru_stresses,
                                                    eq_stress=input_stru_stress,
                                                    vasp=False)
    
        elastic_tensor = result.voigt.tolist()
        bv, gv = result.k_voigt, result.g_voigt
        ev = 9 * bv * gv / (3 * bv + gv)
        uV = (3 * bv - 2 * gv) / (6 * bv + 2 * gv)

        return {
            "elastic_tensor": elastic_tensor,
            "bulk_modulus": float(bv),
            "shear_modulus": float(gv),
            "young_modulus": float(ev),
            "poisson_ratio": float(uV)
        }    
            
    def get_metrics(self, jobs):
        metrics = {}
        for sub_dir in jobs:
            if not os.path.isdir(sub_dir):
                raise ValueError(f"Subdirectory {sub_dir}.")
            
            r = RESULT(path=sub_dir, fmt="abacus")
            metrics[sub_dir] = {i: r[i] for i in ["normal_end", "converge", "scf_steps", "relax_steps","relax_converge","denergy_last", "drho_last", "energy", "stress"]}

        return metrics      
        
        
        
            
            
