import os

import numpy as np

import tbplas as tb

from model import make_cell, make_graphene_sc, save_dat


def test_haydock():
    # Model
    sample = make_graphene_sc(t=-2.7, a=0.142, dim=(4096, 4096, 1))
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "haydock"
    config = tb.Config()
    config.LDOS['recursion_depth'] = 2000
    solver = tb.Solver(sample, config, prefix=prefix)

    # Correlation
    timer = tb.Timer()
    timer.tic("ldos_haydock")
    eng, ldos = solver.calc_ldos_haydock()
    timer.toc("ldos_haydock")
    timer.report_total_time()

    # Output
    save_dat("ldos_py.dat", eng, ldos)


def test_dos():
    # Model
    sample = make_graphene_sc(t=-2.7, a=0.142, dim=(4096, 4096, 1))
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "dos"
    config = tb.Config()
    config.generic['nr_random_samples'] = 1
    config.generic['nr_time_steps'] = 1024
    solver = tb.Solver(sample, config, prefix=prefix)
    analyzer = tb.Analyzer(sample, config, dimension=2)

    # Correlation
    from_scratch = True
    timer = tb.Timer()
    timer.tic("corr_dos")
    if from_scratch:
        corr_dos = solver.calc_corr_dos()
    else:
        corr_dos = np.load(f"sim_data/{prefix}.corr_DOS.npy")
    timer.toc("corr_dos")
    timer.report_total_time()

    # Output
    eng, dos = analyzer.calc_dos(corr_dos)
    save_dat("dos_py.dat", eng, dos)


def test_ac_cond():
    # Model
    # Reference: tbplas/tests/test_units/test_ac.py
    t, a = 3.0, 0.142
    sample = make_graphene_sc(t=t, a=a, dim=(4096, 4096, 1))
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "ac_cond"
    config = tb.Config()
    config.generic['nr_random_samples'] = 1
    config.generic['nr_time_steps'] = 1024
    config.generic['correct_spin'] = True
    solver = tb.Solver(sample, config, prefix=prefix)
    analyzer = tb.Analyzer(sample, config, dimension=2)

    # Correlation
    from_scratch = True
    timer = tb.Timer()
    timer.tic("corr_ac_cond")
    if from_scratch:
        corr = solver.calc_corr_ac_cond()
    else:
        corr = np.load(f"sim_data/{prefix}.corr_AC.npy")
    timer.toc("corr_ac_cond")
    timer.report_total_time()

    # Output
    energies, ac_cond = analyzer.calc_ac_cond(corr)
    save_dat("ac_py.dat", energies/t, ac_cond[0].real*4)


def test_dyn_pol():
    # Model
    # Reference: tbplas/tests/test_units/test_dp.py and test_eps.py
    t, a = 3.0, 0.142
    sample = make_graphene_sc(t=t, a=a, dim=(4096, 4096, 1))
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "dyn_pol"
    config = tb.Config()
    config.generic['nr_random_samples'] = 1
    config.generic['nr_time_steps'] = 1024
    config.generic['correct_spin'] = False
    config.dyn_pol['q_points'] = np.zeros((2, 3), dtype=np.float64)
    config.dyn_pol['q_points'][0] = np.array([[0.86602540, 0.5, 0.0]]) / 0.142
    config.dyn_pol['q_points'][1] = np.array([[4.122280922013927, 2.38, 0.0]])
    solver = tb.Solver(sample, config, prefix=prefix)
    analyzer = tb.Analyzer(sample, config, dimension=2)

    # Correlation
    from_scratch = True
    timer = tb.Timer()
    timer.tic("corr_dyn_pol")
    if from_scratch:
        corr = solver.calc_corr_dyn_pol()
    else:
        corr = np.load(f"sim_data/{prefix}.corr_dyn_pol.npy")
    timer.toc("corr_dyn_pol")
    timer.report_total_time()

    # Output
    q, omegas, dp = analyzer.calc_dyn_pol(corr)
    epsilon = analyzer.calc_epsilon(dp)
    save_dat("dp_py.dat", omegas/t, -dp[0].imag*t*a**2)
    save_dat("eps_py.dat", omegas, epsilon[1].real)


def test_dc_cond():
    # Model
    # Update to C++ version
    sample = make_graphene_sc(t=-2.7, a=0.142, dim=(1024, 1024, 1))
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "dc_cond"
    config = tb.Config()
    config.generic['nr_random_samples'] = 1
    config.generic['nr_time_steps'] = 1024
    config.generic['correct_spin'] = False
    config.DC_conductivity['energy_limits'] = (-2.5, 2.5)
    config.DC_conductivity['nr_time_steps'] = 1024
    solver = tb.Solver(sample, config, prefix=prefix)
    analyzer = tb.Analyzer(sample, config, dimension=2)

    # Correlation
    from_scratch = True
    timer = tb.Timer()
    timer.tic("corr_dc_cond")
    if from_scratch:
        corr_dos, corr_dc = solver.calc_corr_dc_cond()
    else:
        corr_dos = np.load(f"sim_data/{prefix}.corr_DOS.npy")
        corr_dc = np.load(f"sim_data/{prefix}.corr_DC.npy")
    timer.toc("corr_dc_cond")
    timer.report_total_time()

    # Output
    energy, dc = analyzer.calc_dc_cond(corr_dos, corr_dc)
    save_dat("dc_xx_py.dat", energy, dc[0])


def test_hall_cond():
    # Model
    # Reference: tbplas/tests/test_units/test_hall.py
    sample = make_graphene_sc(t=3.0, a=0.142, dim=(1024, 1024, 1))
    sample.set_magnetic_field(100.0)
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "hall_cond"
    config = tb.Config()
    config.generic['nr_random_samples'] = 1
    config.generic['nr_time_steps'] = 1024
    config.generic['correct_spin'] = True
    config.set_temperature(10.0)
    config.dckb['direction'] = 2
    config.dckb['energies'] = np.linspace(-1.0, 1.0, 1000)
    solver = tb.Solver(sample, config, prefix=prefix)
    analyzer = tb.Analyzer(sample, config, dimension=2)

    # Correlation
    from_scratch = True
    timer = tb.Timer()
    timer.tic("hall_cond")
    if from_scratch:
        mu_hall = solver.calc_hall_mu()
    else:
        mu_hall = np.load(f"sim_data/{prefix}.hall_mu.npy")
    timer.toc("hall_cond")
    timer.report_total_time()

    # Output
    omegas, hall = analyzer.calc_hall_cond(mu_hall, unit="h")
    save_dat("hall_cond_py.dat", omegas, hall)


def test_qe():
    # Model
    prim_cell = make_cell(t=-2.7, a=0.142)
    super_cell = tb.SuperCell(prim_cell, dim=(17, 17, 1), pbc=(True, True, False))
    super_cell.add_vacancies([(8, 8, 0, 0)])
    sample = tb.Sample(super_cell)
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "quasi_eigen"
    config = tb.Config()
    config.generic['nr_random_samples'] = 1
    config.generic['nr_time_steps'] = 1024
    config.quasi_eigenstates['energies'] = [-1.0, 0.0, 1.0]
    solver = tb.Solver(sample, config, prefix=prefix)

    # Quasi-eigen states
    timer = tb.Timer()
    timer.tic("qe")
    qs = solver.calc_quasi_eigenstates()
    timer.toc("qe")
    timer.report_total_time()
    np.savetxt("qe_py.dat", qs.T)


def test_wft():
    sample = make_graphene_sc(t=-2.7, a=0.142, dim=(17, 17, 1))
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "wft"
    config = tb.Config()
    config.generic['nr_random_samples'] = 1
    config.generic['nr_time_steps'] = 1024
    solver = tb.Solver(sample, config, prefix=prefix)

    # wft
    timer = tb.Timer()
    timer.tic("wft")
    wf0 = np.zeros((sample.num_orb, 1), dtype=np.complex128)
    for i in range(sample.num_orb):
        if i % 2 == 0:
            wf0[i] = 1.0
    wft = solver.calc_psi_t(wf0, time_log=np.array([0, 1, 1024]))
    timer.toc("wft")
    timer.report_total_time()
    wft = wft.T
    num_row, num_col = wft.shape
    data = np.zeros((num_row, 2*num_col))
    data[:, 0::2] = wft.real
    data[:, 1::2] = wft.imag
    np.savetxt("wft_py.dat", data)


def test_qe_speed():
    # Model
    prim_cell = make_cell(t=-2.7, a=0.142)
    super_cell = tb.SuperCell(prim_cell, dim=(4096, 4096, 1), pbc=(True, True, False))
    sample = tb.Sample(super_cell)
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "quasi_eigen"
    config = tb.Config()
    config.generic['nr_random_samples'] = 1
    config.generic['nr_time_steps'] = 1024
    config.quasi_eigenstates['energies'] = [-1.0, 0.0, 1.0]
    solver = tb.Solver(sample, config, prefix=prefix)

    # Quasi-eigen states
    timer = tb.Timer()
    timer.tic("qe")
    solver.calc_quasi_eigenstates()
    timer.toc("qe")
    timer.report_total_time()


def test_wft_speed():
    sample = make_graphene_sc(t=-2.7, a=0.142, dim=(4096, 4096, 1))
    sample.rescale_ham(9.0)

    # Parameters
    prefix = "wft"
    config = tb.Config()
    config.generic['nr_random_samples'] = 1
    config.generic['nr_time_steps'] = 1024
    solver = tb.Solver(sample, config, prefix=prefix)

    # wft
    timer = tb.Timer()
    timer.tic("wft")
    wf0 = np.zeros((sample.num_orb, 1), dtype=np.complex128)
    for i in range(sample.num_orb):
        if i % 2 == 0:
            wf0[i] = 1.0
    solver.calc_psi_t(wf0, time_log=np.array([0, 1, 1024]))
    timer.toc("wft")
    timer.report_total_time()


def test_bands():
    sample = make_graphene_sc(t=-2.7, a=0.142, dim=(3, 3, 1))
    sample.rescale_ham(9.0)

    k_points = np.array([
        [0.0, 0.0, 0.0],
        [2. / 3, 1. / 3, 0.0],
        [1. / 2, 0.0, 0.0],
        [0.0, 0.0, 0.0],
    ])
    k_path, k_idx = tb.gen_kpath(k_points, [40, 40, 40])
    k_len = tb.gen_kdist(sample.sc0.sc_lat_vec, k_path)

    # Parameters
    config = tb.Config()
    config.generic['nr_time_steps'] = 1024
    config.generic['nr_random_samples'] = 1
    solver = tb.Solver(sample, config)
    analyzer = tb.Analyzer(sample, config, dimension=2)
    energy_all,  dos_all = [], []
    for i_k, kpt in enumerate(k_path):
        sample.reset_array()
        sample.set_k_point(kpt)
        sample.rescale_ham()
        corr_dos = solver.calc_corr_dos()
        eng, dos = analyzer.calc_dos(corr_dos)
        energy_all.append(eng)
        dos_all.append(dos)
    energy_all = np.array(energy_all)
    dos_all = np.array(dos_all)
    if solver.is_master:
        data_dir = "demo_tbpm"
        os.system(f"rm -rf {data_dir} && mkdir {data_dir}")
        np.save(f"{data_dir}/energy", energy_all)
        np.save(f"{data_dir}/dos", dos_all)
        np.save(f"{data_dir}/k_len", k_len)
        np.save(f"{data_dir}/k_idx", k_idx)
