#!/usr/bin/env python 
# -*- coding: utf-8 -*-
# @Time    : 2024/10/18 13:26
# @Author  : 兵
# @email    : 1747193328@qq.com

import os
import traceback


from pathlib import Path

import numpy as np
from PySide6.QtCore import QObject
from loguru import logger

from NepTrainKit.core import MessageManager, Structure, Config
from .base import NepPlotData, StructureData
from .utils import read_nep_out_file, read_atom_num_from_xyz, check_fullbatch, read_nep_in, parse_array_by_atomnum
from ..calculator import process_calculate, process_get_descriptors, _calculate, _get_descriptors, Nep3Calculator, \
    process_get_polarizability


def pca(X, k):
    # 1. 标准化数据（去均值和方差标准化）

    mean = np.mean(X, axis=0)
    X_centered = X - mean


    # 2. 计算协方差矩阵
    cov_matrix = np.cov(X_centered.T)

    # 3. 特征值分解协方差矩阵
    eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)

    # 4. 对特征值进行排序，选择前k个特征值和对应的特征向量
    sorted_indices = np.argsort(eigenvalues)[::-1]  # 从大到小排序
    top_k_eigenvectors = eigenvectors[:, sorted_indices[:k]]

    # 5. 投影到前k个主成分
    X_pca = X_centered.dot(top_k_eigenvectors)

    return X_pca


class ResultData(QObject):
    def __init__(self,nep_txt_path,data_xyz_path,descriptor_path):
        super().__init__()
        structures = Structure.read_multiple(data_xyz_path)
        self.descriptor_path=descriptor_path
        self.data_xyz_path=data_xyz_path
        self.nep_txt_path=nep_txt_path

        self._atoms_dataset=StructureData(structures)
        self.select_index=set()

    @property
    def dataset(self):
        return []

    @property
    def descriptor(self):
        return self._descriptor_dataset

    @property
    def num(self):
        return self._atoms_dataset.num
    @property
    def structure(self):
        return self._atoms_dataset

    def is_select(self,i):
        return i in self.select_index

    def select(self,_list):
        if isinstance(_list,(int,np.int_,np.int64)):
            _list=[_list]

        for i in _list:
            self.select_index.add(i)

    def uncheck(self,_list):
        if isinstance(_list,int):
            _list=[_list]
        for i in _list:
            if i in self.select_index:
                self.select_index.remove(i)



    def export_model_xyz(self,save_path):
        try:
            # row_model=  ase_read(self.data_xyz_path,index=":",format="extxyz")


            # result=[row_model[i] for i in self._atoms_dataset.now_data ]
            # remove=[row_model[i] for i in self._atoms_dataset.remove_data ]
            with open(Path(save_path).joinpath("export_good_model.xyz"),"w",encoding="utf8") as f:
                for structure in self._atoms_dataset.now_data:
                    structure.write(f)

            with open(Path(save_path).joinpath("export_remove_model.xyz"),"w",encoding="utf8") as f:
                for structure in self._atoms_dataset.remove_data:
                    structure.write(f)

            # ase_write(Path(save_path).joinpath("export_good_model.xyz") ,result,append=True)
            # ase_write(Path(save_path).joinpath("export_remove_model.xyz") ,remove,append=True)
            MessageManager.send_info_message(f"File exported to: {save_path}")
        except:
            MessageManager.send_info_message(f"An unknown error occurred while saving. The error message has been output to the log!")
            logger.error(traceback.format_exc())


    def get_atoms(self,index ):

        index=self._atoms_dataset.convert_index(index)
        return self._atoms_dataset.now_data[index][0]



    def remove(self,i):


        self._atoms_dataset.remove(i)
        for dataset in self.dataset:
            dataset.remove(i)
    @property
    def is_revoke(self):
        return self._atoms_dataset.remove_data.size!=0
    def revoke(self):
        self._atoms_dataset.revoke()
        for dataset in self.dataset:
            dataset.revoke( )

    def delete_selected(self ):

        self.remove(list(self.select_index))
        self.select_index.clear()

    def _load_descriptors(self):
        atoms_num_list=[len(structure) for structure in self.structure.now_data]

        if os.path.exists(self.descriptor_path):
            desc_array = read_nep_out_file(self.descriptor_path)

        else:
            desc_array = np.array([])

        if desc_array.size == 0:

            desc_array = process_get_descriptors(self.nep_txt_path.as_posix(), self.structure.now_data)
            if desc_array.size != 0:
                np.savetxt(self.descriptor_path, desc_array, fmt='%.6g')
        else:
            if desc_array.shape[0] != len(atoms_num_list):
                # 原子描述符 需要计算结构描述符
                desc_array = parse_array_by_atomnum(desc_array, atoms_num_list, map_func=np.mean, axis=0)

                pass
            else:
                # 结构描述符
                pass
        if desc_array.size != 0:
            if desc_array.shape[1] > 2:
                desc_array = pca(desc_array, 2)

        self._descriptor_dataset = NepPlotData(desc_array, title="descriptor")


class NepTrainResultData(ResultData):
    def __init__(self,
                 nep_txt_path,
                 data_xyz_path,
                 energy_out_path,
                 force_out_path,
                 stress_out_path,
                 virial_out_path,
                 descriptor_path

                 ):
        super().__init__(nep_txt_path,data_xyz_path,descriptor_path)
        self.energy_out_path = energy_out_path
        self.force_out_path = force_out_path
        self.stress_out_path = stress_out_path
        self.virial_out_path = virial_out_path




        self._load_dataset()
        self._load_descriptors()






    @property
    def dataset(self):
        # return [self.energy, self.stress,self.virial, self.descriptor]
        return [self.energy,self.force,self.stress,self.virial, self.descriptor]



    @property
    def energy(self):
        return self._energy_dataset

    @property
    def force(self):
        return self._force_dataset

    @property
    def stress(self):
        return self._stress_dataset

    @property
    def virial(self):
        return self._virial_dataset



    @classmethod
    def from_path(cls, path,model="train"):
        path = Path(path)
        dataset_path = path.joinpath(f"{model}.xyz")
        if not dataset_path.exists():
            MessageManager.send_message_box(f"{model}.xyz not found in the current working directory.")
            return None
        nep_txt_path = path.joinpath(f"nep.txt")
        if not nep_txt_path.exists():
            MessageManager.send_warning_message(f"nep.txt not found in the current working directory.")


        energy_out_path = path.joinpath(f"energy_{model}.out")
        force_out_path = path.joinpath(f"force_{model}.out")
        stress_out_path = path.joinpath(f"stress_{model}.out")
        virial_out_path = path.joinpath(f"virial_{model}.out")
        descriptor_path = path.joinpath(f"descriptor.out")

        return cls(nep_txt_path,dataset_path,energy_out_path,force_out_path,stress_out_path,virial_out_path,descriptor_path)


    def _load_dataset(self ):
        nep_in = read_nep_in(self.nep_txt_path.with_name("nep.in"))
        atoms_num_list=[len(structure) for structure in self.structure.now_data]
        if (not check_fullbatch(nep_in, len(atoms_num_list))
                or all([
                    not self.energy_out_path.exists(),
                    not self.force_out_path.exists(),
                    not self.stress_out_path.exists(),
                    not self.virial_out_path.exists()

                ])):
            nep_potentials_array, nep_forces_array, nep_virials_array = process_calculate(self.nep_txt_path.as_posix(),
                                                                                          self.structure.now_data)
            try:
                energy_array = np.column_stack(
                    [nep_potentials_array / atoms_num_list, [structure.per_atom_energy for structure in self.structure.now_data]])
                np.savetxt(self.energy_out_path, energy_array, fmt='%10.5f')

            except:
                pass
                energy_array = np.column_stack(
                    [nep_potentials_array / atoms_num_list, nep_potentials_array / atoms_num_list])

            try:

                forces_array = np.column_stack([nep_forces_array,
                                                np.vstack([structure.forces for structure in self.structure.now_data]),

                                                ])
                np.savetxt(self.force_out_path, forces_array, fmt='%10.5f')

            except:
                logger.debug(traceback.format_exc())

                forces_array = np.column_stack([nep_forces_array,
                                                nep_forces_array

                                                ])
            coefficient = (atoms_num_list / np.array([structure.volume for structure in self.structure.now_data]))[:, np.newaxis]

            try:
                virials_array = np.column_stack([nep_virials_array,
                                                 np.vstack([structure.nep_virial for structure in self.structure.now_data]),

                                                 ])

                stress_array = virials_array * coefficient * 160.21766208
                np.savetxt(self.virial_out_path, virials_array, fmt='%10.5f')
                np.savetxt(self.stress_out_path, stress_array, fmt='%10.5f')

            except:
                logger.debug(traceback.format_exc())

                virials_array = np.column_stack([nep_virials_array,
                                                 nep_virials_array

                                                 ])

                stress_array = virials_array * coefficient * 160.21766208

            # MessageManager.send_message_box("Detected that the current mode is not full batch. Please make predictions first, then load!")
            # raise ValueError("Detected that the current mode is not full batch. Please make predictions first, then load!")
        else:
            energy_array = read_nep_out_file(self.energy_out_path)
            forces_array = read_nep_out_file(self.force_out_path)
            virials_array = read_nep_out_file(self.virial_out_path)
            stress_array = read_nep_out_file(self.stress_out_path)
        default_forces = Config.get("widget", "forces_data", "Row")
        if forces_array.size != 0 and default_forces == "Norm":

            forces_array = parse_array_by_atomnum(forces_array, atoms_num_list, map_func=np.linalg.norm, axis=0)

            self._force_dataset = NepPlotData(forces_array, title="force")
        else:
            self._force_dataset = NepPlotData(forces_array, group_list=atoms_num_list, title="force")

        self._energy_dataset = NepPlotData(energy_array, title="energy")

        if float(nep_in.get("lambda_v", 1)) != 0:
            self._stress_dataset = NepPlotData(stress_array, title="stress")

            self._virial_dataset = NepPlotData(virials_array, title="virial")
        else:
            self._stress_dataset = NepPlotData([], title="stress")

            self._virial_dataset = NepPlotData([], title="virial")



class NepPolarizabilityResultData(ResultData):
    def __init__(self,
                 nep_txt_path,
                 data_xyz_path,
                 polarizability_out_path,

        descriptor_path
                 ):
        super().__init__(nep_txt_path,data_xyz_path,descriptor_path)


        self.polarizability_out_path = polarizability_out_path


        self._load_dataset()

        self._load_descriptors()

        self._atoms_polarizability_diagonal_dataset = NepPlotData([], title="")
        self._atoms_polarizability_no_diagonal_dataset = NepPlotData([], title="")


    @property
    def dataset(self):

        return [self.polarizability_diagonal,self.polarizability_no_diagonal,self.atoms_polarizability_diagonal,self.atoms_polarizability_no_diagonal, self.descriptor]

    @property
    def num(self):
        return self._atoms_dataset.num
    @property
    def structure(self):
        return self._atoms_dataset

    @property
    def polarizability_diagonal(self):
        return self._polarizability_diagonal_dataset
    @property
    def polarizability_no_diagonal(self):
        return self._polarizability_no_diagonal_dataset

    @property
    def atoms_polarizability_diagonal(self):
        return self._atoms_polarizability_diagonal_dataset
    @property
    def atoms_polarizability_no_diagonal(self):
        return self._atoms_polarizability_no_diagonal_dataset

    @property
    def descriptor(self):
        return self._descriptor_dataset



    @classmethod
    def from_path(cls, path,model="train"):
        path = Path(path)
        dataset_path = path.joinpath(f"{model}.xyz")
        if not dataset_path.exists():
            MessageManager.send_message_box(f"{model}.xyz not found in the current working directory.")
            return None
        nep_txt_path = path.joinpath(f"nep.txt")
        if not nep_txt_path.exists():
            MessageManager.send_message_box(f"nep.txt not found in the current working directory.")
            return None

        polarizability_out_path = path.joinpath(f"polarizability_{model}.out")

        descriptor_path = path.joinpath(f"descriptor.out")

        return cls(nep_txt_path, dataset_path, polarizability_out_path, descriptor_path)
    def _load_dataset(self ):
        nep_in = read_nep_in(self.nep_txt_path.with_name("nep.in"))
        atoms_num_list=np.array([len(structure) for structure in self.structure.now_data])
        if (not check_fullbatch(nep_in, len(atoms_num_list))
                or all([
                    not self.polarizability_out_path.exists(),


                ])):

            try:
                nep_polarizability_array = process_get_polarizability(self.nep_txt_path.as_posix(),self.structure.now_data)

                nep_polarizability_array=nep_polarizability_array/(atoms_num_list[:, np.newaxis])
                polarizability_array = np.column_stack([nep_polarizability_array,
                                                np.vstack([structure.nep_polarizability for structure in self.structure.now_data]),

                                                ])


                np.savetxt(self.polarizability_out_path, polarizability_array, fmt='%10.5f')

            except:
                logger.debug(traceback.format_exc())

                polarizability_array = np.column_stack([nep_polarizability_array,
                                                nep_polarizability_array

                                                ])




        else:
            polarizability_array = read_nep_out_file(self.polarizability_out_path)



        self._polarizability_diagonal_dataset = NepPlotData(polarizability_array[:, [0,1,2,6,7,8]], title="Polar Diag")

        self._polarizability_no_diagonal_dataset = NepPlotData(polarizability_array[:, [3,4,5,9,10,11]], title="Polar NoDiag")
