# -*- coding: utf-8 -*-
from pathlib import Path

from baderkit.core import Bader, Grid

from simmate.apps.baderkit.models.bader import Bader as BaderModel
from simmate.workflows.core import Workflow


class Bader__Baderkit__Bader(Workflow):
    required_files = ["AECCAR0", "AECCAR2", "CHGCAR", "POTCAR"]
    use_database = True
    database_table = BaderModel
    use_previous_directory = ["AECCAR0", "AECCAR2", "CHGCAR", "POTCAR"]
    # parent_workflows = [
    #     "population-analysis.vasp-baderkit.bader-warren-lab",
    #     ]

    """
    Runs a Bader charge analysis on VASP outputs using the BaderKit package.
    """

    @classmethod
    def run_config(
        cls,
        previous_directory: Path,
        source: dict = None,
        directory: Path = None,
        run_id=None,
        **kwargs,
    ):
        # create CHGCAR_sum grid
        grid1 = Grid.from_vasp(directory / "AECCAR0")
        grid2 = Grid.from_vasp(directory / "AECCAR2")
        reference_grid = grid1.linear_add(grid2)
        # load CHGCAR
        charge_grid = Grid.from_vasp(directory / "CHGCAR")
        # create Bader
        bader = Bader(
            charge_grid=charge_grid,
            reference_grid=reference_grid,
            **kwargs,
        )
        # get the table for this workflow and update entry
        datatable = cls.database_table.objects.get(run_id=run_id)
        datatable.update_from_baderkit(bader, directory)
