# -*- coding: utf-8 -*-

import math


class Kpoints:
    """
    This class takes a structure along with a KptGrid or KptPath object and
    then generates a VASP formatted KPOINTS file for you. Alternatively,
    you can provide a structre and a KPOINTS file to get back a KptGrid or
    KptPath object. This behaves as a converter class, just like the POSCAR class.
    """

    @staticmethod
    def write_file(
        # structure to base off of. this will be passed through the kpoint_method
        structure,
        # either a KptGrid or KptPath object. TODO - KptGrid is just a float (density)
        kpoint_method,  # !!! JUST A FLOAT VALUE FOR NOW
        filename="KPOINTS",
        comment=None,
    ):

        # let's have a list where we keep track of list of lines that we
        # will join at the end
        lines = []

        # first line is just a comment. If the user didn't supply one, I will
        # just use "Generated by Simmate with [name of kpoint_method]"
        if not comment:
            comment = f"Generated by Simmate at {str(kpoint_method)} kpts per Ang^-1"
        lines.append(comment)

        # check to see if kpoint_method is a KptGrid object
        # TODO - for now this is just a float value as I didn't write the class
        # NOTE - while we suport this, the prefered method is to set kpt_density
        # through the KSPACING tag in the INCAR.
        # BUG - make sure integers are converted to floats, even though a kpt
        # density of 1 or greater should never be used, we should still be
        # ready for it.
        if isinstance(kpoint_method, float):

            # append 0 to the second line which tells VASP that the grid should
            # be generated automatically based on the lines below it
            lines.append("0")

            # The next line indicates whether we use a Gamma-centered or a
            # Monkhorst-Pack grid. We can determine this based off of what
            # type of lattice the structure has:
            #   Gamma centered for hexagonal cells
            #   Otherwise use a Monkhorst-Pack grid
            # TODO -- should this be inside the KptGrid class? I not sure if
            # this is universal or unique to VASP.
            if structure.lattice.is_hexagonal():
                lines.append("Gamma")
            else:
                lines.append("Monkhorst-Pack")

            # Now write out the kpt grid in the form "N1 N2 N3"
            # TODO - this code should be implemented in the KptGrid class
            # with a method that I then reformat. This also includes whether
            # this is a singleline or not as well as if there is a shift
            # applied at the end
            # iterate through the reciprocal lattice vectors and find the values
            # of N to use for each.
            n_list = []
            for vector in structure.lattice.reciprocal_lattice.abc:
                # check how many subdivisions we need to cut this reciprocal
                # lattice vector into in order to get X kpts per Angstrom^-1
                # I want to round up to the nearst integer, which requires
                # the math.ceil function (as opposed to the "//" operator)
                # TODO - reciprocal_density = kpoint_method until I add KptGrid class
                nsubdivisions = math.ceil(vector / kpoint_method)
                n_list.append(nsubdivisions)
            # now combine the list and add it
            lines.append(" ".join(str(n) for n in n_list))

            # TODO - if there is a specified shift in the kpts, add them
            # if kpoint_method.shift:
            #     lines.append(kpoint_method.shift)

        # TODO - I haven't written the KptPath object yet so I can't do this.
        # chek to see if kpoint_method is a KptPath object
        # elif isinstance(kpoint_method, KptPath):
        #     pass

        # Now we can join all of the lines!
        lines = "\n".join(lines) + "\n"

        # Now write the file
        # If I want a separate to_string method, I should move only this code.
        with open(filename, "w") as file:
            file.write(lines)

    # TODO -- I have no reason to read kpt files yet. I likely won't implement
    # this until I add the KptGrid and KptPath classes as well as start doing
    # Band structure calculations
    # @staticmethod
    # def from_file():
    #     pass
