import csv
import math

import numpy as np

import atom
import crystal
import density
import eigen


class Settings:
    pass


# Set up the object "s" that contains all settings for the run.
# After this function, "s" is not modified. Adjust settings here.
def initialize():
    s = Settings()

    # Resolution: store functions on an NxNxN grid
    s.Nres = 16

    # Lattice constant of conventional cell, in Bohr radii
    #a = 6.7463 # C
    a = 10.261 # Si
    #a = 10.298 # GaP
    #a = 10.683 # Ge/GaAs
    #a = 11.449 # InAs

    # Crystal structure, choose Cubic{Simple,BodyCentered,FaceCentered}
    s.crystal = crystal.CubicFaceCentered(s.Nres, a)
    # Grid dimensions for k-point sampling of Brillouin zone
    s.bzmesh = (3, 3, 3)

    # Sequence of Brillouin zone symmetry points for band structure plot
    s.bzpath = "LGXWULWXUG"
    # Number of substeps to take between two symmetry points in bzpath
    s.Nbzsub = 10

    # Number of orbitals to use per atom (but not necessarily converge)
    s.Norb = 8
    # Number of total Kohn-Sham eigenstates to converge
    s.Neig = 12

    # Elements and positions of atoms in the unit cell.
    # Positions are expressed in the cell basis (a1,a2,a3).
    atoms = [
        ("Si", (-0.125, -0.125, -0.125)),
        ("Si", ( 0.125,  0.125,  0.125)),
    ]

    # Plane-wave cutoff energy in Hartree
    s.Ecut = 50

    # Finalize the "s" object. Avoid editing below this line.

    # Get the grid point coordinates in direct and Fourier space
    s.rx, s.ry, s.rz, s.dr = s.crystal.get_grid_direct()
    s.qx, s.qy, s.qz, s.dq = s.crystal.get_grid_fourier()
    # Remember to call fftshift before plotting!
    s.qx = np.fft.ifftshift(s.qx)
    s.qy = np.fft.ifftshift(s.qy)
    s.qz = np.fft.ifftshift(s.qz)

    # For convenience
    s.r2 = s.rx**2 + s.ry**2 + s.rz**2
    s.q2 = s.qx**2 + s.qy**2 + s.qz**2

    # Place the atoms on the grid as specified in "atoms"
    s.atoms = []
    basis = s.crystal.a
    for a in atoms:
        pos = a[1][0] * basis[0] + a[1][1] * basis[1] + a[1][2] * basis[2]
        s.atoms.append(atom.Atom(s, a[0], pos))

    return s


def main():
    # Get settings object "s"
    s = initialize()

    # It would be madness to try to converge the electron density
    # for all band structure k-points at the same time, so instead
    # we converge it for a small set and then assume it's fixed.
    rho, Veff = density.get_selfconsistent_density(s)
    print("\nCalculating band structure with fixed electron density")

    # Get list of positions of high-symmetry k-points
    ksteps = s.crystal.get_brillouin_zone_path(s.bzpath)
    Nsteps = len(ksteps) - 1

    # Total distance in k-space, used for first CSV column
    kdistance = 0.0

    # Open output CSV file for writing
    file = open("result.csv", "w", newline="")
    csvw = csv.writer(file, quoting=csv.QUOTE_MINIMAL)

    # Loop over symmetry points, except last one
    for i in range(Nsteps):
        print("Band calculation part {}/{} running...".format(i + 1, Nsteps))

        # Subdivide path between this point and next point
        if i < Nsteps - 1:
            ksubs = np.linspace(ksteps[i], ksteps[i + 1], s.Nbzsub, endpoint=False)
        else: # include endpoint in last iteration
            ksubs = np.linspace(ksteps[i], ksteps[i + 1], s.Nbzsub + 1, endpoint=True)

        # Loop over points on subdivided path
        for j, kp in enumerate(ksubs):
            # Solve Kohn-Sham equation for this k-point, with fixed Veff
            pf = "    ({}/{}) ".format(j + 1, len(ksubs))
            energies, psis_r, psis_q = eigen.solve_kohnsham_equation(s, Veff, kp, pf)

            csvw.writerow([kdistance] + energies[0 : s.Neig])

            # Keep track of how far we moved in k-space
            kd = np.linalg.norm(ksteps[i + 1] - ksteps[i])
            kdistance += kd / s.Nbzsub

    file.close()


if __name__ == "__main__":
    main()