import matplotlib.pyplot as plt
import export_lattice as ex
import numpy as np
import pybinding as pb

from pybinding.repository.graphene import t, a_cc, a

# define lattice of monolayer graphene with a_cc [nm] interatomic distance and t [eV] hopping,
# EnergyScale is the scaling factor of the hopping parameters, important for the rescaling of the spectral quantity.
energy_scale = 9.5


def graphene(onsite=(0, 0)):
    """Return the basic lattice specification for monolayer graphene with nearest neighbor"""

    from math import sqrt

    # create a lattice with 2 primitive vectors
    lat = pb.Lattice(a1=[a, 0], a2=[a/2, a/2 * sqrt(3)])

    # Add sublattices
    lat.add_sublattices(
        # name, position, and onsite potential
        ('A', [0, -a_cc/2], onsite[0]),
        ('B', [0, +a_cc/2], onsite[1])
    )

    # Add hoppings
    lat.add_hoppings(
        # inside the main cell, between which atoms, and the value
        ([0, 0], 'A', 'B', t / energy_scale),
        # between neighboring cells, between which atoms, and the value
        ([1, -1], 'A', 'B', t / energy_scale),
        ([0, -1], 'A', 'B', t / energy_scale),
    )

    return lat


lattice, disorder, disorded_structural = graphene()

# number of decomposition parts in each direction of matrix. This divides the lattice into various sections, each of which is calculated in parallel
nx = ny = 1

# number of unit cells in each direction.
lx = 256
ly = 256

# make config object which caries info about
# - the number of decomposition parts [nx, ny],
# - lengths of structure [lx, ly]
# - boundary conditions, setting True as periodic boundary conditions, and False elsewise,
# - info if the exported hopping and onsite data should be complex,
# - info of the precision of the exported hopping and onsite data, 0 - float, 1 - double, and 2 - long double.
configuration = ex.Configuration(divisions=[nx, ny], length=[lx, ly], boundaries=[True, True],
                                 is_complex=False, precision=1, energy_scale=energy_scale)

# make calculation object which caries info about
# - the name of the function
#   DOS - denstity of states == 1,
#   CondXX - conductivity in xx direction == 2,
#   CondXY - conductivity in xy direction == 3,
#   OptCond - optical conductivity == 4
#   SpinCond - spin conductivity == 5
#   SingleCondXX - single energy XX conductivity == 6
#   SingleCondXY - single energy XY conductivity == 7
# - number of moments for the calculation,
# - number of different random vector realisations,
# - number of disorder realisations.
# - energy and gamma for single energy calculations.
calculation = ex.Calculation(fname=['DOS'], num_moments=[500], num_random=[1], num_disorder=[1])

# make modification object which caries info about (TODO: Other modifications can be added here)
# - magnetic field can be set to True. Default case is False. In exported file it's converted to 1 and 0.
modification = ex.Modification(magnetic_field=False)

# export the lattice from the lattice object, config and calculation object and the name of the file
# the disorder is optional. If there is disorder in the lattice for now it should be given separately
ex.export_lattice(lattice, configuration, calculation, modification, 'graphene_example.h5')

# plotting the lattice
lattice.plot()
plt.show()