import concurrent.futures
from copy import copy, deepcopy
from datetime import datetime, timedelta
from enum import Enum
import os
from os import PathLike
from pathlib import Path
from typing import Union

from nemspy import ModelingSystem
import numpy as np

from coupledmodeldriver import Platform
from coupledmodeldriver.generate.schism.configure import (
    # NEMSSCHISMRunConfiguration,
    SCHISMRunConfiguration,
)
from coupledmodeldriver.generate.schism.script import (
    SchismEnsembleRunScript,
    SchismEnsembleCleanupScript,
    SchismRunJob,
    SchismCombineHotstartJob,
    SchismCombineSchoutJob,
)
from coupledmodeldriver.script import SlurmEmailType
from coupledmodeldriver.utilities import (
    create_symlink,
    get_logger,
    ProcessPoolExecutorStackTraced,
)

LOGGER = get_logger('cplmdldrv')


class RunPhase(Enum):
    """
    options for run phase of SCHISM execuction
    """

    COLDSTART = 'coldstart'
    HOTSTART = 'hotstart'


def generate_schism_configuration(
    configuration_directory: PathLike,
    output_directory: PathLike = None,
    relative_paths: bool = False,
    overwrite: bool = False,
    parallel: bool = True,
):
    """
    generate an SCHISM run configuration from JSON configuration files generated by ``initialize_schism``

    :param configuration_directory: path containing JSON configuration files
    :param output_directory: path to store generated configuration files
    :param relative_paths: whether to write relative paths in generated configuration files
    :param overwrite: whether to overwrite existing files
    :param parallel: generate configurations concurrently
    """

    start_time = datetime.now()

    if not isinstance(configuration_directory, Path):
        configuration_directory = Path(configuration_directory)

    if output_directory is None:
        output_directory = configuration_directory
    elif not isinstance(output_directory, Path):
        output_directory = Path(output_directory)

    output_directory = output_directory.resolve()
    if not output_directory.is_absolute():
        output_directory = output_directory.absolute()

    if not output_directory.exists():
        os.makedirs(output_directory, exist_ok=True)

    if configuration_directory.absolute().resolve() != Path.cwd():
        starting_directory = Path.cwd()
        os.chdir(configuration_directory)
        configuration_directory = Path.cwd()
    else:
        starting_directory = None

    use_nems = 'configure_nems.json' in [
        filename.name.lower() for filename in configuration_directory.iterdir()
    ]

    if use_nems:
        LOGGER.debug(f'generating NEMS configuration')
        pass
        # TODO
        # base_configuration = NEMSSCHISMRunConfiguration.read_directory(configuration_directory)
    else:
        LOGGER.debug(f'generating SCHISM-only configuration')
        base_configuration = SCHISMRunConfiguration.read_directory(configuration_directory)

    base_configuration.move_paths(configuration_directory)

    platform = base_configuration['modeldriver']['platform']

    job_duration = base_configuration['slurm']['job_duration']
    slurm_account = base_configuration['slurm']['account']
    partition = base_configuration['slurm']['partition']
    email_type = base_configuration['slurm']['email_type']
    email_address = base_configuration['slurm']['email_address']

    # TODO update functions for SCHISM
    original_fgrid_filename = base_configuration['schism']['fgrid_path']
    original_hgrid_filename = base_configuration['schism']['hgrid_path']
    schism_processors = base_configuration['schism']['processors']
    spinup_duration = base_configuration['schism']['tidal_spinup_duration']
    use_original_mesh = base_configuration['schism']['use_original_mesh']

    if original_hgrid_filename is None or not original_hgrid_filename.exists():
        raise FileNotFoundError(f'horizontal grid not found at "{original_hgrid_filename}"')

    local_fgrid_filename = output_directory / original_fgrid_filename.name
    local_hgrid_filename = output_directory / 'hgrid.gr3'
    local_param_filename = output_directory / 'param.nml'

    do_spinup = spinup_duration is not None

    run_phase = 'HOTSTART' if do_spinup else 'COLDSTART'

    if slurm_account is None:
        slurm_account = platform.value['slurm_account']

    ensemble_run_script_filename = output_directory / f'run_{platform.name.lower()}.sh'
    ensemble_cleanup_script_filename = output_directory / f'cleanup.sh'

    if use_original_mesh:
        LOGGER.info(
            f'using original mesh from "{os.path.relpath(original_hgrid_filename.resolve(), Path.cwd())}"'
        )
        if original_fgrid_filename.exists():
            create_symlink(original_fgrid_filename, local_fgrid_filename)
        create_symlink(original_hgrid_filename, local_hgrid_filename)
    elif overwrite or not local_hgrid_filename.exists():
        LOGGER.info(
            f'rewriting original mesh to "{os.path.relpath(local_hgrid_filename.resolve(), Path.cwd())}"'
        )
        try:
            # TODO: COPY MESH FILE
            ...
        except Exception as error:
            LOGGER.warning(error)

    if local_param_filename.exists():
        os.remove(local_param_filename)

    if local_fgrid_filename.exists():
        base_configuration['schism']['fgrid_path'] = local_fgrid_filename
    if local_hgrid_filename.exists():
        base_configuration['schism']['hgrid_path'] = local_hgrid_filename

    runs_directory = output_directory / 'runs'
    if not runs_directory.exists():
        runs_directory.mkdir(parents=True, exist_ok=True)

    perturbations = base_configuration.perturb()

    LOGGER.info(
        f'generating {len(perturbations)} run configuration(s) in "{os.path.relpath(runs_directory.resolve(), Path.cwd())}"'
    )

    if parallel:
        process_pool = ProcessPoolExecutorStackTraced()
        LOGGER.info(f'leveraging {os.cpu_count()} processor(s)')
    else:
        process_pool = None

    futures = []

    if do_spinup:
        spinup_directory = output_directory / 'spinup'

        spinup_configuration = copy(base_configuration)

        if not spinup_configuration.files_exist(spinup_directory):
            spinup_kwargs = {
                'directory': spinup_directory,
                'configuration': spinup_configuration,
                'duration': spinup_duration,
                'local_fgrid_filename': local_fgrid_filename,
                'local_hgrid_filename': local_hgrid_filename,
                'link_mesh': True,
                'relative_paths': relative_paths,
                'overwrite': overwrite,
                'platform': platform,
                'schism_processors': schism_processors,
                'slurm_account': slurm_account,
                'job_duration': job_duration,
                'partition': partition,
                'email_type': email_type,
                'email_address': email_address,
            }

            if parallel:
                futures.append(process_pool.submit(write_spinup_directory, **spinup_kwargs))
            else:
                spinup_directory = write_spinup_directory(**spinup_kwargs)
                LOGGER.info(f'wrote configuration to "{spinup_directory}"')
    else:
        spinup_directory = None

    for run_name, run_configuration in perturbations.items():
        run_directory = runs_directory / run_name

        if not run_configuration.files_exist(run_directory):
            run_kwargs = {
                'directory': run_directory,
                'name': run_name,
                'phase': run_phase,
                'configuration': run_configuration,
                'local_fgrid_filename': local_fgrid_filename,
                'local_hgrid_filename': local_hgrid_filename,
                'link_mesh': use_original_mesh,
                'relative_paths': relative_paths,
                'overwrite': overwrite,
                'platform': platform,
                'schism_processors': schism_processors,
                'slurm_account': slurm_account,
                'job_duration': job_duration,
                'partition': partition,
                'email_type': email_type,
                'email_address': email_address,
                'do_spinup': do_spinup,
                'spinup_directory': spinup_directory,
            }

            if parallel:
                futures.append(process_pool.submit(write_run_directory, **run_kwargs))
            else:
                write_run_directory(**run_kwargs)
                LOGGER.info(f'wrote configuration to "{run_directory}"')

    if parallel:
        for completed_future in concurrent.futures.as_completed(futures):
            LOGGER.info(f'wrote configuration to "{completed_future.result()}"')

    cleanup_script = SchismEnsembleCleanupScript()
    LOGGER.debug(
        f'writing cleanup script "{os.path.relpath(ensemble_cleanup_script_filename.resolve(), Path.cwd())}"'
    )
    cleanup_script.write(filename=ensemble_cleanup_script_filename, overwrite=overwrite)

    LOGGER.info(
        f'writing ensemble run script "{os.path.relpath(ensemble_run_script_filename.resolve(), Path.cwd())}"'
    )
    run_job_script = SchismEnsembleRunScript(
        platform=platform,
        commands=[
            'echo deleting previous SCHISM output',
            f'sh {ensemble_cleanup_script_filename.name}',
            'echo deleting previous SCHISM logs',
            'rm spinup/*.log runs/*/*.log',
        ],
        run_spinup=do_spinup,
    )
    run_job_script.write(ensemble_run_script_filename, overwrite=overwrite)

    if starting_directory is not None:
        os.chdir(starting_directory)

    LOGGER.info(f'finished in {datetime.now() - start_time}')

    if starting_directory is not None:
        os.chdir(starting_directory)

    LOGGER.info(f'finished in {datetime.now() - start_time}')


def write_spinup_directory(
    directory: PathLike,
    # configuration: Union[SCHISMRunConfiguration, NEMSSCHISMRunConfiguration],
    configuration: SCHISMRunConfiguration,
    duration: timedelta,
    local_hgrid_filename: PathLike,
    local_fgrid_filename: PathLike = None,
    link_mesh: bool = False,
    relative_paths: bool = False,
    overwrite: bool = False,
    platform: Platform = None,
    schism_processors: int = None,
    slurm_account: str = None,
    job_duration: timedelta = None,
    partition: str = None,
    email_type: SlurmEmailType = None,
    email_address: str = None,
) -> Path:
    """
    write an SCHISM tidal spinup configuration to the given directory

    :param directory: path to output directory
    :param configuration: run configuration
    :param duration: tidal spinup duration
    :param relative_paths: whether to write relative paths (instead of absolute)
    :param overwrite: whether to overwrite existing files
    :param platform: HPC platform to configure for
    :param schism_processors: number of processors to assign to SCHISM execution
    :param slurm_account: Slurm account name
    :param job_duration: wall-clock duration of Slurm job
    :param partition: Slurm partition name
    :param email_type: Slurm email notification option
    :param email_address: Slurm email notification address
    """

    if not isinstance(directory, Path):
        directory = Path(directory)

    if not directory.exists():
        directory.mkdir(parents=True, exist_ok=True)

    if not overwrite and configuration.files_exist(directory):
        LOGGER.warning(f'skipping existing directory "{directory}"')
        return directory

    job_name = 'SCHISM_COLDSTART_SPINUP'
    combine_job_name = 'SCHISM_COMBINE_HOTSTARTFILES_SPINUP'

    pyschism_driver = configuration.pyschism_driver

    start_date = configuration['schism']['modeled_start_time']
    end_date = configuration['schism']['modeled_end_time']
    spinup_duration = configuration['schism']['tidal_spinup_duration']

    pyschism_driver.param.opt.start_date = start_date - spinup_duration
    pyschism_driver.param.core.rnday = spinup_duration
    pyschism_driver.param.core.ihfskip = spinup_duration
    pyschism_driver.param.schout.nhot = 1
    pyschism_driver.param.schout.nhot_write = pyschism_driver.param.core.ihfskip

    if relative_paths:
        configuration.relative_to(directory, inplace=True)

    schism_old_io = False
    if 'nems' in configuration:
        # TODO: Check if nems uses the new or old io run
        nems = configuration.nemspy_modeling_system
        nems = ModelingSystem(
            nems.start_time - duration,
            nems.start_time,
            nems.interval,
            ocn=deepcopy(nems['OCN']),
            **nems.attributes,
        )
        processors = nems.processors
        model_executable = configuration['nems']['executable_path']
        hotstart_combiner = configuration['schism']['schism_hotstart_combiner_path']
    else:
        nems = None
        processors = schism_processors
        model_executable = configuration['schism']['schism_executable_path']
        schism_old_io = configuration['schism']['schism_use_old_io']
        hotstart_combiner = configuration['schism']['schism_hotstart_combiner_path']

    source_filename = configuration['schism']['source_filename']

    model_executable = update_path_relative(model_executable, relative_paths, directory)
    hotstart_combiner = update_path_relative(
        hotstart_combiner, relative_paths, directory / 'outputs'
    )
    source_filename = update_path_relative(source_filename, relative_paths, directory)

    job_script_filename = directory / 'schism.job'
    job_script = SchismRunJob(
        platform=platform,
        old_io=schism_old_io,
        slurm_tasks=processors,
        slurm_account=slurm_account,
        slurm_duration=job_duration,
        slurm_run_name=job_name,
        executable=model_executable,
        slurm_partition=partition,
        slurm_email_type=email_type,
        slurm_email_address=email_address,
        slurm_error_filename=f'{job_name}.err.log',
        slurm_log_filename=f'{job_name}.out.log',
        source_filename=source_filename,
    )
    job_script.write(job_script_filename, overwrite=overwrite)

    combine_script_filename = directory / 'combine_hotstart.job'
    combine_script = SchismCombineHotstartJob(
        platform=platform,
        iterations_idx=[pyschism_driver.param.schout.nhot_write],
        slurm_tasks=processors,
        slurm_account=slurm_account,
        slurm_duration=job_duration,
        slurm_run_name=combine_job_name,
        executable=hotstart_combiner,
        slurm_partition=partition,
        slurm_email_type=email_type,
        slurm_email_address=email_address,
        slurm_error_filename=f'{combine_job_name}.err.log',
        slurm_log_filename=f'{combine_job_name}.out.log',
        source_filename=source_filename,
    )
    combine_script.write(combine_script_filename, overwrite=overwrite)

    # Link to the combined hotstart file
    try:
        combined_hotstart_filename = (
            f'hotstart_it={pyschism_driver.param.schout.nhot_write}.nc'
        )
        create_symlink(
            directory / 'outputs' / combined_hotstart_filename,
            directory / 'hotstart_out.nc',
            relative=True,
        )
    except:
        LOGGER.warning(f'unable to link `{combined_hotstart_filename}`')

    if 'nems' in configuration:
        LOGGER.debug(f'setting spinup to {duration}')

        nems.write(
            directory, overwrite=overwrite, include_version=True,
        )
        LOGGER.info(
            f'writing NEMS+SCHISM tidal spinup configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
        )
    else:
        LOGGER.debug(
            f'writing SCHISM tidal spinup configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
        )

    pyschism_driver.write(directory, overwrite=overwrite)

    return directory


def write_run_directory(
    directory: PathLike,
    name: str,
    phase: str,
    # configuration: Union[SCHISMRunConfiguration, NEMSSCHISMRunConfiguration],
    configuration: SCHISMRunConfiguration,
    local_hgrid_filename: PathLike,
    local_fgrid_filename: PathLike = None,
    link_mesh: bool = False,
    relative_paths: bool = False,
    overwrite: bool = False,
    platform: Platform = None,
    schism_processors: int = None,
    slurm_account: str = None,
    job_duration: timedelta = None,
    partition: str = None,
    email_type: SlurmEmailType = None,
    email_address: str = None,
    do_spinup: bool = False,
    spinup_directory: PathLike = None,
) -> Path:
    """
    write an SCHISM tidal spinup configuration to the given directory

    :param directory: path to output directory
    :param name: run name
    :param phase: SCHISM run phase
    :param configuration: run configuration
    :param relative_paths: whether to write relative paths (instead of absolute)
    :param overwrite: whether to overwrite existing files
    :param platform: HPC platform to configure for
    :param schism_processors: number of processors to assign to SCHISM execution
    :param slurm_account: Slurm account name
    :param job_duration: wall-clock duration of Slurm job
    :param partition: Slurm partition name
    :param email_type: Slurm email notification option
    :param email_address: Slurm email notification address
    :param do_spinup: whether to use spinup results to hotstart run
    :param spinup_directory: path to spinup directory
    """

    if not isinstance(directory, Path):
        directory = Path(directory)
    if spinup_directory is not None and not isinstance(spinup_directory, Path):
        spinup_directory = Path(spinup_directory)

    if not directory.exists():
        directory.mkdir(parents=True, exist_ok=True)
    LOGGER.debug(
        f'writing run configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
    )

    if not overwrite and configuration.files_exist(directory):
        LOGGER.warning(f'skipping existing directory "{directory}"')
        return directory

    job_name = f'SCHISM_{phase}_{name}'
    combine_job_name = f'SCHISM_COMBINE_HOTSTARTFILES_{name}'
    combine2_job_name = f'SCHISM_COMBINE_OUTPUTFILES_{name}'

    pyschism_driver = configuration.pyschism_driver

    pyschism_driver.param.opt.dramp = None
    pyschism_driver.param.opt.drampbc = None
    pyschism_driver.param.opt.dramp_ss = None
    pyschism_driver.param.opt.drampwind = None
    if do_spinup:
        pyschism_driver.param.opt.ihot = 1

    if relative_paths:
        configuration.relative_to(directory, inplace=True)

    schism_old_io = False
    if 'nems' in configuration:
        # TODO: Check if nems uses the new or old io run
        nems = configuration.nemspy_modeling_system
        processors = nems.processors
        model_executable = configuration['nems']['executable_path']
        hotstart_combiner = configuration['schism']['schism_hotstart_combiner_path']
        output_combiner = configuration['schism']['schism_schout_combiner_path']
    else:
        nems = None
        processors = schism_processors
        model_executable = configuration['schism']['schism_executable_path']
        schism_old_io = configuration['schism']['schism_use_old_io']
        hotstart_combiner = configuration['schism']['schism_hotstart_combiner_path']
        output_combiner = configuration['schism']['schism_schout_combiner_path']

    source_filename = configuration['schism']['source_filename']

    model_executable = update_path_relative(model_executable, relative_paths, directory)
    hotstart_combiner = update_path_relative(
        hotstart_combiner, relative_paths, directory / 'outputs'
    )
    output_combiner = update_path_relative(
        output_combiner, relative_paths, directory / 'outputs'
    )
    source_filename = update_path_relative(source_filename, relative_paths, directory)

    job_script_filename = directory / 'schism.job'
    job_script = SchismRunJob(
        platform=platform,
        old_io=schism_old_io,
        slurm_tasks=processors,
        slurm_account=slurm_account,
        slurm_duration=job_duration,
        slurm_run_name=job_name,
        executable=model_executable,
        slurm_partition=partition,
        slurm_email_type=email_type,
        slurm_email_address=email_address,
        slurm_error_filename=f'{job_name}.err.log',
        slurm_log_filename=f'{job_name}.out.log',
        source_filename=source_filename,
    )
    job_script.write(job_script_filename, overwrite=overwrite)

    rnday = pyschism_driver.param.core.rnday
    ihfskip = pyschism_driver.param.core.ihfskip
    nhot = pyschism_driver.param.schout.nhot
    nhot_write = pyschism_driver.param.schout.nhot_write
    dt = pyschism_driver.param.core.dt

    if nhot == 1 and nhot_write is not None:

        # num_hs_files = int(nhot_write / ihfskip)
        num_hs_files = int(round(timedelta(days=rnday).total_seconds() / (dt * nhot_write)))
        hotstart_iters_idx = [nhot_write * (mult + 1) for mult in range(num_hs_files)]

        combine_script_filename = directory / 'combine_hotstart.job'
        combine_script = SchismCombineHotstartJob(
            platform=platform,
            iterations_idx=hotstart_iters_idx,
            slurm_tasks=processors,
            slurm_account=slurm_account,
            slurm_duration=job_duration,
            slurm_run_name=combine_job_name,
            executable=hotstart_combiner,
            slurm_partition=partition,
            slurm_email_type=email_type,
            slurm_email_address=email_address,
            slurm_error_filename=f'{combine_job_name}.err.log',
            slurm_log_filename=f'{combine_job_name}.out.log',
            source_filename=source_filename,
        )
        combine_script.write(combine_script_filename, overwrite=overwrite)

    if schism_old_io:
        end_output_stack_idx = int(
            np.ceil(np.ceil(timedelta(days=rnday).total_seconds() / dt) / ihfskip)
        )

        combine2_script_filename = directory / 'combine_output.job'
        combine2_script = SchismCombineSchoutJob(
            platform=platform,
            start_idx=1,
            end_idx=end_output_stack_idx,
            slurm_tasks=processors,
            slurm_account=slurm_account,
            slurm_duration=job_duration,
            slurm_run_name=combine2_job_name,
            executable=output_combiner,
            slurm_partition=partition,
            slurm_email_type=email_type,
            slurm_email_address=email_address,
            slurm_error_filename=f'{combine2_job_name}.err.log',
            slurm_log_filename=f'{combine2_job_name}.out.log',
            source_filename=source_filename,
        )
        combine2_script.write(combine2_script_filename, overwrite=overwrite)

    if 'nems' in configuration:
        nems.write(
            directory, overwrite=overwrite, include_version=True,
        )
        LOGGER.info(
            f'writing NEMS+SCHISM run configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
        )
    else:
        LOGGER.debug(
            f'writing SCHISM run configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
        )

    pyschism_driver.write(
        directory, overwrite=overwrite,
    )

    if do_spinup:
        try:
            create_symlink(
                spinup_directory / 'hotstart_out.nc', directory / 'hotstart.nc', relative=True,
            )
        except:
            LOGGER.warning(
                f'unable to link hotstart from coldstart to hotstart; '
                'you must manually link or copy this file after coldstart completes'
            )

    return directory


def update_path_relative(
    path: PathLike, relative: bool = False, relative_directory: PathLike = None
) -> Path:
    """
    update the given path based on the given relative path

    :param path: input path
    :param relative: whether to use a relative path to the input path
    :param relative_directory: directory to use to update input directory
    :return: updated path
    """

    if path is not None:
        if not isinstance(path, Path):
            path = Path(path)
        if relative_directory is None:
            relative_directory = Path.cwd()
        if relative:
            if not isinstance(relative_directory, Path):
                relative_directory = Path(relative_directory)
            if path.is_absolute():
                path = Path(os.path.relpath(path, relative_directory))
        elif not path.is_absolute():
            path = (relative_directory / path).resolve().absolute()
    return path
