from ai2_kit.core.artifact import Artifact, ArtifactDict
from ai2_kit.core.script import BashTemplate
from ai2_kit.core.script import BashScript, BashStep
from ai2_kit.core.job import JobFuture, gather_jobs
from ai2_kit.core.log import get_logger
from ai2_kit.core.util import dict_nested_get, expand_globs

from pydantic import BaseModel
from typing import List, Tuple
from dataclasses import dataclass
import os
import copy
import random
import sys
import json

from .iface import ICllTrainOutput, BaseCllContext
from .data import DataFormat, get_data_format

from .constant import (
    DP_CHECKPOINT_FILE,
    DP_DISP_FILE,
    DP_PROFILING_FILE,
    DP_INPUT_FILE,
    DP_FROZEN_MODEL,
    DP_ORIGINAL_MODEL
)

logger = get_logger(__name__)


class CllDeepmdInputConfig(BaseModel):

    model_num: int = 4
    """
    Total number of models to train.
    """
    init_dataset: List[str]
    """
    Dataset used to initialize training.
    """
    input_template: dict = dict()
    """
    Deepmd input template.
    """
    compress_model: bool = False
    """
    Whether to compress model after training.
    """
    isolate_outliers: bool = False
    """
    If isolate_outliers is enabled, then outlier data will be separated from training data.
    """
    outlier_f_cutoff: float = 10.
    """
    The threshold of force magnitude to determine whether a data is outlier.
    """
    outlier_weight: float = 0.003
    """
    The weight of outlier data in training data.
    """

    fixture_models: List[str] = []
    """
    Fixture models used to initialize training, support glob pattern.
    If this is not empty, then the whole training process will be skipped.
    This feature is useful for debugging, or explore more structures without training.
    The models should be on the remote executor.
    The name fixture is used as the concept of fixture in pytest.
    """

    group_by_formula: bool = False
    """
    Grouping dataset by formula
    If this is enabled, then the dataset will be grouped by formula.
    Otherwise, the dataset will be grouped by ancestor.

    Set this to True when you have multiple structures with the same ancestor.
    """


class CllDeepmdContextConfig(BaseModel):
    script_template: BashTemplate
    dp_cmd: str = 'dp'


@dataclass
class CllDeepmdInput:
    config: CllDeepmdInputConfig
    type_map: List[str]
    old_dataset: List[Artifact]  # training data used by previous iteration
    new_dataset: List[Artifact]  # training data used by current iteration


@dataclass
class CllDeepmdContext(BaseCllContext):
    config: CllDeepmdContextConfig


@dataclass
class GenericDeepmdOutput(ICllTrainOutput):
    models: List[Artifact]
    dataset: List[Artifact]

    def get_mlp_models(self) -> List[Artifact]:
        return self.models

    def get_training_dataset(self) -> List[Artifact]:
        return self.dataset


async def cll_deepmd(input: CllDeepmdInput, ctx: CllDeepmdContext):
    executor = ctx.resource_manager.default_executor

    # use fixture models
    if len(input.config.fixture_models) > 0:
        logger.info(f'Using fixture models: {input.config.fixture_models}')
        model_paths = executor.run_python_fn(expand_globs)(input.config.fixture_models)
        assert len(model_paths) > 0, f'No fixture models found: {input.config.fixture_models}'
        return GenericDeepmdOutput(
            dataset=input.old_dataset,
            models=[Artifact.of(url=url, format=DataFormat.DEEPMD_MODEL)
                    for url in model_paths]
        )

    # setup workspace
    work_dir = os.path.join(executor.work_dir, ctx.path_prefix)
    [new_dataset_dir, tasks_dir, outlier_dir] = executor.setup_workspace(
        work_dir, ['new_dataset', 'tasks', 'outlier_dataset'])

    init_dataset = ctx.resource_manager.resolve_artifacts(input.config.init_dataset)
    # input dataset contains data that generated by previous iteration
    # it should not contain the initial dataset
    input_dataset: List[Artifact] = input.old_dataset

    # make new dataset from raw data
    new_dataset, outlier_dataset = executor.run_python_fn(make_deepmd_dataset)(
        dataset_dir=new_dataset_dir,
        outlier_dir=outlier_dir,
        raw_data_collection=[a.to_dict() for a in input.new_dataset],
        type_map=input.type_map,
        isolate_outliers=input.config.isolate_outliers,
        outlier_f_cutoff=input.config.outlier_f_cutoff,
        group_by_formula=input.config.group_by_formula,
    )

    input_dataset += [ Artifact.of(**a) for a in new_dataset]
    # use attrs to distinguish outlier dataset
    input_dataset += [ Artifact.of(**{**a, 'attrs': {**a['attrs'], 'outlier': True}}) for a in outlier_dataset]

    #TODO: separate job creation and submission
    jobs: List[JobFuture] = []
    output_dirs = []

    # train multiple models at the same time
    for i in range(input.config.model_num):
        # each model should be trained in its own task_dir
        task_dir = os.path.join(tasks_dir, str(i).zfill(3))
        executor.mkdir(task_dir)

        # create dp train input file
        # NOTE: dp v1 format is supported currently
        # TODO: support more params if it is necessary

        # ref: https://github.com/deepmodeling/dpgen2/blob/master/examples/ch4/param_CH4_deepmd-kit-2.1.1.json
        # ref: https://github.com/deepmodeling/dpgen2/blob/master/dpgen2/op/prep_dp_train.py
        # ref: https://github.com/deepmodeling/dpgen2/blob/master/dpgen2/op/run_dp_train.py

        dp_input = copy.deepcopy(input.config.input_template)
        training: dict = dp_input['training']

        # set output files
        training['disp_file'] = DP_DISP_FILE
        training['save_ckpt'] = DP_CHECKPOINT_FILE
        training['profiling_file'] = DP_PROFILING_FILE

        # set random seed
        discriptor = dp_input['model']['descriptor']
        if discriptor['type'] == 'hybrid':
            for d in discriptor['list']:
                d['seed'] = _random_seed()
        else:
            discriptor['seed'] = _random_seed()
        dp_input['model']['fitting_net']['seed'] = _random_seed()
        dp_input['training']['seed'] = _random_seed()

        # set training data
        train_systems: List[str] = []
        outlier_systems: List[str] = []
        validation_systems: List[str] = []

        for a in input_dataset + init_dataset:
            if dict_nested_get(a.attrs, ['deepmd', 'validation_data'], False):
                validation_systems.append(a.url)
            else:
                if dict_nested_get(a.attrs, ['outlier'], False):
                    outlier_systems.append(a.url)
                else:
                    train_systems.append(a.url)


        auto_prob_str = "prob_sys_size"
        if input.config.isolate_outliers:
            # if isolate_outlier is enabled, then we need to set different weight for outlier data
            # e.g: prob_sys_size;0:57:0.9997;57:60:0.0003
            auto_prob_str += f';0:{len(train_systems)}:{1-input.config.outlier_weight}'
            auto_prob_str += f';{len(train_systems)}:{len(train_systems)+len(outlier_systems)}:{input.config.outlier_weight}'

        # Drop v1 support
        # TODO: remove in the future
        # training['systems'] = train_systems + outlier_systems
        # training.setdefault('set_prefix', 'set')  # respect user input
        # training.setdefault('batch_size', 'auto')
        # training['auto_prob_style'] = auto_prob_str

        # v2 training data
        training_data = training.get('training_data', {})
        training_data = {
            'set_prefix': 'set',
            'batch_size': 'auto',
            'sys_probs': None,
            ** training_data,
            'systems': train_systems + outlier_systems,
            'auto_prob_style': auto_prob_str,
        }
        training['training_data'] = training_data

        # ignore validation section if no data is provided, or else dp will throw error
        # OSError: [Errno cannot find valid a data system] Please check your setting for data systems
        if len(validation_systems) > 0:
            validation_data = {
                'systems': validation_systems,
                'set_prefix': training_data['set_prefix'],
                'batch_size': training_data['batch_size'],
            }
            training['validation_data'] = validation_data

        # other params
        dp_input['model']['type_map'] = input.type_map

        # write config to executor
        dp_input_text = json.dumps(dp_input, indent=2)
        dp_input_path = os.path.join(task_dir, DP_INPUT_FILE)
        executor.dump_text(dp_input_text, dp_input_path)

        # build script
        dp_cmd = ctx.config.dp_cmd
        dp_train_cmd = [dp_cmd, 'train', DP_INPUT_FILE]

        steps = [
            BashStep(cmd=dp_train_cmd, checkpoint='dp-train') # type: ignore
        ]

        if input.config.compress_model:
            dp_freeze_cmd = [dp_cmd, 'freeze', '-o', DP_ORIGINAL_MODEL]
            dp_compress_cmd = \
                [dp_cmd, 'compress', '-i', DP_ORIGINAL_MODEL, '-o', DP_FROZEN_MODEL]
            steps.append(BashStep(cmd=dp_freeze_cmd)) # type: ignore
            steps.append(BashStep(cmd=dp_compress_cmd)) # type: ignore
        else:
            dp_freeze_cmd = [dp_cmd, 'freeze', '-o', DP_FROZEN_MODEL]
            steps.append(BashStep(cmd=dp_freeze_cmd)) # type: ignore

        dp_train_script = BashScript(
            template=ctx.config.script_template,
            steps=steps # type: ignore
        )
        output_dirs.append(task_dir)

        # submit job
        job = executor.submit(dp_train_script.render(), cwd=task_dir, checkpoint_key=f'queue-job:dp-train:{task_dir}:{i}')
        jobs.append(job)

    await gather_jobs(jobs, max_tries=2)

    logger.info(f'All models are trained, output dirs: {output_dirs}')
    return GenericDeepmdOutput(
        dataset=input_dataset.copy(),
        models=[Artifact.of(
            url=os.path.join(url, DP_FROZEN_MODEL),
            format=DataFormat.DEEPMD_MODEL,
        ) for url in output_dirs]
    )


def _random_seed():
    return random.randrange(sys.maxsize) % (1 << 32)


def __export_remote_functions():
    import dpdata
    from itertools import groupby


    def make_deepmd_dataset(
        dataset_dir: str,
        outlier_dir: str,
        raw_data_collection: List[ArtifactDict],
        isolate_outliers: bool,
        outlier_f_cutoff: float,
        type_map: List[str],
        group_by_formula: bool = False,
    ):
        dataset_collection: List[Tuple[ArtifactDict, dpdata.LabeledSystem]] = []
        outlier_collection: List[Tuple[ArtifactDict, dpdata.LabeledSystem]] = []

        for raw_data in raw_data_collection:
            data_format = get_data_format(raw_data)  # type: ignore
            dp_system = None
            if data_format == DataFormat.CP2K_OUTPUT_DIR:
                dp_system = dpdata.LabeledSystem(os.path.join(raw_data['url'], 'output'), fmt='cp2k/output', type_map=type_map)
            elif data_format == DataFormat.VASP_OUTPUT_DIR:
                dp_system = dpdata.LabeledSystem(os.path.join(raw_data['url'], 'OUTCAR'), fmt='vasp/outcar', type_map=type_map)
            else:
                raise ValueError(f"Unsupported data format: {data_format}")
            # one case of len(dp_system) == 0 is when the system is not converged
            if dp_system is None or 0 == len(dp_system):
                continue  # skip invalid data

            if isolate_outliers and dp_system.data['forces'].max() > outlier_f_cutoff:
                outlier_collection.append((raw_data, dp_system))
            else:
                dataset_collection.append((raw_data, dp_system))

        _write_dp_dataset = _write_dp_dataset_by_formula if group_by_formula else _write_dp_dataset_by_ancestor

        dataset_dirs = _write_dp_dataset(dp_system_list=dataset_collection, out_dir=dataset_dir, type_map=type_map)
        outlier_dirs = _write_dp_dataset(dp_system_list=outlier_collection, out_dir=outlier_dir, type_map=type_map)

        return dataset_dirs, outlier_dirs


    def _write_dp_dataset_by_formula(dp_system_list: List[Tuple[ArtifactDict, dpdata.LabeledSystem]], out_dir: str, type_map: List[str]):
        """
        Write dp dataset that grouping by formula
        Use dpdata.MultipleSystems to merge systems with the same formula
        Use this when group by ancestor not works for you
        """
        if len(dp_system_list) == 0:
            return []
        os.makedirs(out_dir, exist_ok=True)

        multi_systems = dpdata.MultiSystems(dp_system_list[0][1])
        for _, system in dp_system_list[1:]:
            multi_systems.append(system)

        multi_systems.to_deepmd_npy(out_dir, type_map=type_map)  # type: ignore

        return [ {
            'url': os.path.join(out_dir, fname),
            'format': DataFormat.DEEPMD_NPY,
            'attrs': {},  # it is meaning less to set attrs in this case
        } for fname in os.listdir(out_dir)]


    def _write_dp_dataset_by_ancestor(dp_system_list: List[Tuple[ArtifactDict, dpdata.LabeledSystem]], out_dir: str, type_map: List[str]):
        """
        write dp dataset that grouping by ancestor
        """
        output_dirs: List[ArtifactDict] = []
        for key, dp_system_group in groupby(dp_system_list, key=lambda x: x[0]['attrs']['ancestor']):
            dp_system_group = list(dp_system_group)
            if 0 == len(dp_system_group):
                continue  # skip empty dataset
            group_out_dir = os.path.join(out_dir, key.replace('/', '_'))

            # merge dp_systems with the same ancestor into single data set
            dp_system = dp_system_group[0][1]
            for item in dp_system_group[1:]:
                dp_system += item[1]
            dp_system.to_deepmd_npy(group_out_dir, set_size=len(dp_system), type_map=type_map)  # type: ignore
            # inherit attrs key from input artifact
            output_dirs.append({'url': group_out_dir,
                                'format': DataFormat.DEEPMD_NPY,
                                'attrs': {**dp_system_group[0][0]['attrs']}})  # type: ignore
        return output_dirs


    return (make_deepmd_dataset, )

(
    make_deepmd_dataset,
) = __export_remote_functions()
