from .model import Model
import json, os
import math
import shutil, glob, copy
from pathlib import Path
from typing import List, Tuple

import numpy as np
from ase.vibrations import Vibrations
from ase.thermochemistry import HarmonicThermo

from abacustest import AbacusSTRU, WriteInput, ReadInput
from abacustest.lib_model.comm import copy_abacusjob
from abacustest.constant import RECOMMAND_IMAGE, RECOMMAND_COMMAND, RECOMMAND_MACHINE


class VibrationModel(Model):
    @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 vibration frequency 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('-s', "--stepsize", default=0.01, type=float, help="The step size for finite difference calculation, default is 0.01 Angstrom, which should be a small positive value.")
        parser.add_argument('-i', "--index", default=None, nargs="*", type=int, help="The indices of atoms to calculate the vibration frequency, starts from 1. You can specify multiple indices separated by space. If not specified, all atoms will be considered.")
        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.")
        
        if params.stepsize <= 0:
            raise ValueError("Step size should be a positive value.")

        folders = []
        for i, job_path in enumerate(params.job):
            disped_stru_job_paths = prepare_abacus_vibration_analysis(job_path,
                                                                      params.index,
                                                                      params.stepsize)

            folders.append(disped_stru_job_paths)
        
        setting = {
            "save_path": "results",
            "run_dft": [{
                "ifrun": True,
                "example": folders,
                "command": f"{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 'run.sh', 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 vibration results:\n\tabacustest model vibration 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 vibration frequency 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.')
        parser.add_argument('-t', '--temperature', default=[298.15], nargs="+", type=float, help="Temperature settings for thermochemistry calculations. Can be one or three values: "
        "1. Single temperature: -t 298.15 (calculate thermo properties at 298.15 K only) "
        "2. Temperature range: -t 100 500 41 (calculate thermo properties at 40 evenly spaced temperatures from 100K to 500K. ). "
        "Default is 298.15 K.")
        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.")
        
        if len(params.temperature) == 1:
            temperature = [params.temperature[0]]
        elif len(params.temperature) == 3:
            t_start, t_end, t_num = params.temperature
            t_num = int(t_num)
            temperature = np.linspace(t_start, t_end, t_num).tolist()
        else:
            raise ValueError("Number of arguments for -t or --temperature should be 1 or 3")

        results = {}
        for job_path in params.job:
            result = post_abacus_vibration_analysis_onejob(job_path, temperature)
            results[job_path] = result
        
        # Save the vibration analysis results
        json.dump(results, open("metrics_vibration.json", "w"), indent=4)
        print("\nThe vibration analysis results are saved in 'metrics_vibration.json'")


def prepare_abacus_vibration_analysis(job_path: Path,
                                      selected_atoms: List[int] = None,
                                      stepsize: float = 0.01):
        """
        Prepare ABACUS input files for vibrational analysis.
        """
        input_params = ReadInput(os.path.join(job_path, "INPUT"))
        input_params['calculation'] = 'scf'
        input_params['cal_force'] = 1

        stru_file = input_params.get('stru_file', "STRU")
        stru = AbacusSTRU.read(os.path.join(job_path, stru_file), fmt="stru")
        displaced_stru = copy.deepcopy(stru)
        original_stru_coord = np.array(stru.coords)

        if selected_atoms is None:
            selected_atoms = [i for i in range(stru.natoms)]
        else:
            selected_atoms = [i-1 for i in selected_atoms] # Use 0-based index
            selected_atoms.sort()
        
        vib_work_path = os.path.join(job_path, "vib")

        if os.path.exists(vib_work_path) and os.path.isdir(vib_work_path):
            shutil.rmtree(vib_work_path) # Clear the existed job folder
        
        vib_job_paths = []

        # Prepare ABACUS input files for the given structure
        abacus_scf_work_path = os.path.join(vib_work_path, "SCF")
        original_stru_job_path = os.path.join(abacus_scf_work_path, "eq")
        os.makedirs(original_stru_job_path)

        copy_abacusjob(job_path, original_stru_job_path, input_file=False)
        WriteInput(input_params, os.path.join(original_stru_job_path, "INPUT"))
        vib_job_paths.append(original_stru_job_path)

        # Dump selected atoms and stepsize info to a json file used in postprocess
        prepare_params = {
            "selected_atoms": selected_atoms,
            "stepsize": stepsize
        }
        with open(os.path.join(original_stru_job_path, "prepare_params.json"), "w") as f:
            json.dump(prepare_params, f, indent=2)

        # Prepare ABACUS input files for each displaced structure.
        STEP_MAP = {'+': 1, '-': -1}
        for selected_atom in selected_atoms:
            for dire_idx, direction in enumerate(['x', 'y', 'z']):
                displaced_stru_coord = copy.deepcopy(original_stru_coord)
                for step in STEP_MAP.keys(): # Two steps along one direction
                    disped_stru_job_path = os.path.join(abacus_scf_work_path, f"disp_{selected_atom+1}_{direction}{step}")
                    copy_abacusjob(job_path, disped_stru_job_path, input_file=False, stru=False)

                    displaced_stru_coord[selected_atom][dire_idx] = original_stru_coord[selected_atom][dire_idx] + stepsize * STEP_MAP[step]
                    displaced_stru.coords = displaced_stru_coord
                    WriteInput(input_params, os.path.join(disped_stru_job_path, "INPUT"))
                    displaced_stru.write(os.path.join(disped_stru_job_path, stru_file))

                    vib_job_paths.append(disped_stru_job_path)

        return vib_job_paths

def identify_complex_types(complex_array):
    real_part = np.real(complex_array)
    imag_part = np.imag(complex_array)

    is_real = np.isclose(imag_part, 0)
    is_pure_imag = np.isclose(real_part, 0) & ~np.isclose(imag_part, 0)
    is_general = ~is_real & ~is_pure_imag

    return is_real, is_pure_imag, is_general

def collect_force(abacusjob_dir):
    from abacustest.lib_collectdata.collectdata import RESULT
    
    abacusresult = RESULT(fmt="abacus", path=abacusjob_dir)
    metrics = {i: abacusresult[i] for i in ['force', 'normal_end', 'converge']}
    if metrics['normal_end'] is not True:
        print(f"ABACUS calculation in {abacusjob_dir} didn't end normally")
    elif metrics['converge'] is not True:
        print(f"ABACUS calculation in {abacusjob_dir} didn't reached SCF convergence")
    else:
        pass

    return metrics['force']

def dump_cache_forces_json(stru: AbacusSTRU,
                           work_dir: Path,
                           disped_stru_job_paths: List[str]
):
    """
    Dump cache forces json files for ase.vibration module needed.
    """
    vib_cache_dir = os.path.join(work_dir, "vib/cache")
    os.makedirs(vib_cache_dir, exist_ok=True)
    cache_forces_json = {"forces": {
        "__ndarray__": [[stru.natoms, 3],
                        "float64",
                        []]
    }}
    for disped_stru_job_path in disped_stru_job_paths:
        cache_forces_json["forces"]["__ndarray__"][2] = collect_force(disped_stru_job_path)
        disped_stru_job_basename = os.path.basename(disped_stru_job_path)
        if os.path.basename(disped_stru_job_path) == "eq":
            cache_file = "cache.eq.json"
        else:
            # get filename from job path
            infos = disped_stru_job_basename.split("_")
            atom_index, step = int(infos[1]) - 1, infos[2]
            cache_file = f"cache.{atom_index}{step}.json"
        with open(os.path.join(vib_cache_dir, cache_file), "w") as fin:
            json.dump(cache_forces_json, fin, indent=2)
    
    return vib_cache_dir

def post_abacus_vibration_analysis_onejob(work_dir: Path,
                                          temperature: List[float] = [298.15]):
    """
    Post-process ABACUS vibration analysis results for one job.
    """
    # Read parameters from prepare process
    with open(os.path.join(work_dir, "vib/SCF/eq/prepare_params.json"), "r") as fin:
        prepare_params = json.load(fin)
    
    # Save cached json file as ase.vibration needed
    stru = AbacusSTRU.read(os.path.join(work_dir, "vib/SCF/eq/STRU"), fmt="stru")
    disped_stru_job_paths = glob.glob(os.path.join(work_dir, "vib/SCF/*"))
    vib_cache_dir = dump_cache_forces_json(stru, work_dir, disped_stru_job_paths)
    
    vib = Vibrations(stru.to(fmt='ase'),
                     name=vib_cache_dir,
                     indices=prepare_params['selected_atoms'],
                     delta=prepare_params['stepsize'],
                     nfree=2)

    # Do the vibration analysis
    vib.summary()
    # Generate list of frequencies in the return value
    frequencies = vib.get_frequencies()
    real_freq_mask, imag_freq_mask, complex_freq_mask = identify_complex_types(frequencies)
    real_freq, imag_freq = np.real(frequencies[real_freq_mask]).tolist(), frequencies[imag_freq_mask].tolist()
    for key, value in enumerate(imag_freq):
        imag_freq[key] = -math.fabs(value.imag) # Represent imaginary frequency with negative number
    freqs = imag_freq + real_freq

    # Write animations of normal modes in ASE traj format
    vib.write_mode()

    # Thermochemistry calculations
    # Vibrations.get_energies() gets `h \nu` for each mode, which is from the eigenvalues of force constant
    # matrix. The force constant matrix should be a real symmetric matrix mathematically, but due to numerical
    # errors during calculating its matrix element, it will deviate from symmetric matrix slightly, and its eigenvalue
    # will have quite small imaginary parts. Magnitude of imaginary parts will decrease as accuracy of force
    # increases, and it's safe to use norm of the complex eigenvalue as vibration energy if force is 
    # accurate enough.
    vib_energies = vib.get_energies()
    vib_energies_float = [float(np.linalg.norm(i)) for i in vib_energies]
    zero_point_energy = sum(vib_energies_float) / 2
    thermo = HarmonicThermo(vib_energies, ignore_imag_modes=True)

    thermo_corr = {}
    for t in temperature:
        entropy = thermo.get_entropy(t)
        free_energy = thermo.get_helmholtz_energy(t)
        thermo_corr[f'{t}K'] = {'entropy': float(entropy),
                                'free_energy': float(free_energy)}
    
    return {'frequencies': freqs,
            'zero_point_energy': float(zero_point_energy),
            'thermo_corr': thermo_corr}
