"""Workflow for C2DB.

Important utility functions in this script:

  create_tasks
  get_mag_states
  get_most_stable_magstate
  task
  is_thermodynamically_stable
  is_dynamically_stablex
  all_done

The workflow is factored into four components:

  magstate_workflow
  therm_stability_workflow
  dynamical_stability_workflow
  property_workflow

"""
from myqueue.task import task as mqtask
from pathlib import Path
from asr.core import read_json, chdir
from asr.setup.strains import main as setupstrains
from asr.setup.strains import get_strained_folder_name, get_relevant_strains
from ase.io import read
import os

VERBOSE = os.environ.get('MQVERBOSE', False)


def task(*args, **kwargs):
    """Get MyQueue task instance."""
    name = kwargs.get("name") or args[0]
    if "creates" not in kwargs:
        kwargs["creates"] = [f"results-{name}.json"]
    return mqtask(*args, **kwargs)


class MissingUnrelaxedStructure(Exception):
    """Exception raised when unrelaxed.json is missing."""

    pass


def is_thermodynamically_stable(folder):
    """Determine whether a material is thermodynamically stable."""
    ch = read_json(folder / "results-asr.convex_hull.json")
    hform = ch["hform"]

    if hform > 0.2:
        return False
    return True


def is_dynamically_stable(folder):
    """Determine whether a material is dynamically stable."""
    ph = read_json(folder / "results-asr.phonons.json")
    if ph["dynamic_stability_phonons"] == 'low':
        return False

    st = read_json(folder / "results-asr.stiffness.json")
    if st["dynamic_stability_stiffness"] == 'low':
        return False

    return True


def get_cwd():
    """Get current working directory."""
    return Path('.').absolute()


reference_dbs = [
    '/home/niflheim2/cmr/databases/referencedatabases/oqmd123.db',
    '/home/niflheim2/cmr/C2DB-ASR/collected-databases/references-c2db.db'
]


def basic_workflow(folder):
    """Generate tasks related finding most stable magnetic structures."""
    tasks = []
    if not (Path(folder) / 'structure.json').is_file():
        if not (Path(folder) / 'unrelaxed.json').is_file():
            raise MissingUnrelaxedStructure
        tasks += [task("asr.relax", resources="40:20h",
                       folder=folder),
                  task("asr.structureinfo", resources="1:10m",
                       deps=["asr.relax"], folder=folder)]
    else:
        tasks += [task("asr.structureinfo", resources="1:10m", folder=folder)]

    tasks += [
        task("asr.gs", folder=folder,
             resources="40:5h", deps=["asr.structureinfo"]),
        task("asr.magstate", folder=folder,
             deps=["asr.gs"]),
        task(
            'asr.convex_hull ' + ' '.join(reference_dbs),
            folder=folder,
            name='asr.convex_hull',
            resources='1:10m',
            deps=['asr.gs']),
        task("asr.bandstructure",
             folder=folder,
             resources="40:5h",
             deps=["asr.gs"]),
        task("asr.projected_bandstructure",
             folder=folder,
             resources="1:10m",
             deps=["asr.bandstructure"]),
        task('asr.pdos',
             folder=folder,
             resources='40:5h',
             deps=['asr.gs']),
        task('asr.bader', deps=['asr.gs'],
             folder=folder,
             tmax='1h'),
    ]

    return tasks


def dynamical_stability_workflow(folder):
    """Generate tasks related to determining dynamical stability."""
    verbose_print('Executing dynamical_stability_workflow().')
    if isinstance(folder, str):
        folder = Path(folder)

    tasks = []
    with chdir(folder):
        if not setupstrains.done:
            setupstrains()

    relevant_strains = get_relevant_strains(pbc=[True, True, False])
    strainfolders = [Path(
        get_strained_folder_name(sign * 1, i, j, clamped=False)
    ).resolve() for i, j in relevant_strains for sign in [-1, 1]]

    assert strainfolders, 'Missing strain folders.'
    for strainfolder in strainfolders:
        tasks += [task(
            "asr.relax",
            folder=strainfolder,
            resources="40:3h",
        )]

    tasks += [task(
        "asr.stiffness",
        resources="1:10m",
        folder=folder,
        deps=([f'{strainfolder}/asr.relax' for strainfolder
               in strainfolders])
    )]

    tasks += [task("asr.phonons", folder=folder,
                   restart=3, resources="40:8h",
                   deps=["asr.gs"])]
    return tasks, strainfolders


def all_done(list_of_tasks):
    """Determine if all tasks in list_of_tasks are done."""
    return all([task.check_creates_files()
                for task in list_of_tasks])


def property_workflow(folder, strainfolders):
    """Generate tasks for various material properties."""
    verbose_print('Executing property_workflow()')
    atoms = read(folder / 'structure.json')
    tasks = [
        task("asr.polarizability",
             folder=folder,
             weight=1,
             cores=120,
             tmax='20h',
             deps=["asr.gs"]),
    ]

    gsresults = read_json(folder / "results-asr.gs.json")
    structureinfo = read_json(folder / "results-asr.structureinfo.json")
    gap = gsresults.get("gap")
    gap_dir = gsresults.get("gap_dir")
    gap_nosoc = gsresults["gap_nosoc"]
    topology_file = folder / "results-asr.berry.json"
    if topology_file.is_file():
        topresult = read_json(topology_file)
        topology = topresult['Topology']
    else:
        topology = None

    gap_tolerance = 0.01

    if gap > gap_tolerance:
        tasks += [task("asr.emasses",
                       folder=folder,
                       resources="40:5h",
                       deps=["asr.gs"]),
                  task("asr.emasses@validate",
                       folder=folder,
                       resources="1:10m",
                       deps=["asr.emasses"])]

        for strainfolder in strainfolders:
            tasks += [task(
                "asr.gs",
                folder=strainfolder,
                resources="40:8h",
            )]

        tasks += [
            task("asr.deformationpotentials",
                 resources="40:8h",
                 folder=folder,
                 deps=([f'{strainfolder}/asr.gs' for strainfolder
                        in strainfolders]))]
    else:
        tasks += [task(
            "asr.fermisurface",
            folder=folder,
            resources="1:10m",
            deps=["asr.gs"],
        )]

    if gap_nosoc > gap_tolerance:
        verbose_print(get_cwd(), 'Has band gap.')
        tasks += [task("asr.hse@calculate",
                       folder=folder,
                       restart=2,
                       resources="80:50h")]
        tasks += [task("asr.hse",
                       folder=folder,
                       resources="40:10h",
                       deps=["asr.hse@calculate"])]

        # Topologically non-trivial materials are especially hard.
        # We filter these away in case the material have been found to
        # have a non-trivial topology.
        if topology in {'Trivial', None}:
            tasks += [
                task("asr.borncharges",
                     folder=folder,
                     resources="40:5h",
                     restart=3,
                     deps=["asr.gs"]),
                task("asr.infraredpolarizability",
                     folder=folder,
                     deps=[
                         "asr.borncharges",
                         "asr.phonons",
                         "asr.polarizability",
                     ]),
            ]
            if not structureinfo['has_inversion_symmetry']:
                tasks.append(
                    task("asr.piezoelectrictensor",
                         resources="40:20h",
                         restart=3,
                         deps=["asr.stiffness"]))

        if len(atoms) < 5 and gap > 0.2:
            tasks.append(task("asr.gw", folder=folder, resources="120:48h"))
    else:
        tasks += [
            task(
                "asr.plasmafrequency",
                folder=folder,
                resources="40:20h",
                deps=["asr.gs"],
            )]

    if gap > 0 and gap_dir < 0.7:
        tasks += [
            task("asr.berry",
                 folder=folder,
                 resources="240:xeon40_768:10h",
                 deps=["asr.gs"])]

    magresults = read_json(folder / "results-asr.magstate.json")
    magnetic = magresults.get("is_magnetic")
    if magnetic:
        tasks += [
            task("asr.exchange",
                 folder=folder,
                 resources="40:10h",
                 deps=["asr.gs"])]

        tasks += [
            task("asr.orbmag",
                 folder=folder,
                 resources="1:20m",
                 deps=["asr.magstate"])]
    else:
        if (not structureinfo['has_inversion_symmetry'] and gap_nosoc > 0):
            tasks += [
                task("asr.shift",
                     folder=folder,
                     resources="240:20h",
                     deps=["asr.gs"])]

    return tasks


def return_tasks(tasks):
    """Wrap function for returning tasks."""
    if VERBOSE:
        print(get_cwd(), tasks)
    return tasks


def verbose_print(*args):
    """Only print if VERBOSE."""
    if VERBOSE:
        print(*args)


def create_tasks():
    """Create MyQueue Task list for the C2DB workflow.

    Note that this workflow relies on the folder layout of C2DB so be
    careful.

    """
    verbose_print(get_cwd())
    folder = Path('.')
    tasks = []
    tasks += basic_workflow(folder)

    if not all_done(tasks):
        return tasks

    if not is_thermodynamically_stable(folder):
        verbose_print('Is not thermodynamically stable.')
        return tasks

    atoms = read(folder / 'structure.json')

    if len(atoms) > 20:
        verbose_print('More than 20 atoms in cell.'
                      ' Not calculating dynamic stability.')
        return tasks

    dyn_tasks, strain_folders = dynamical_stability_workflow(folder)
    tasks += dyn_tasks
    if not all_done(tasks):
        verbose_print('Missing dynamic stability tasks.')
        return tasks

    if not is_dynamically_stable(folder):
        verbose_print('Is not dynamically stable.')
        return tasks

    tasks += property_workflow(folder, strain_folders)

    return tasks


if __name__ == '__main__':
    tasks = create_tasks()

    for ptask in tasks:
        print(ptask, ptask.is_done())