#! /usr/bin/env python
"""Script for checking the model using Python version of TBPLaS."""
import argparse
from typing import List

import numpy as np
import tbplas as tb


class FakeCachedSample:
    """Fake CachedSample class for plotting."""
    def init_array(self) -> None:
        pass


def extract(content: List[str], label: str) -> List[List[str]]:
    i0 = content.index(f"begin_{label}\n")
    i1 = content.index(f"end_{label}\n")
    lines = [_.split() for _ in content[i0+1:i1]]
    return lines


def main():
    # Parse command-line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("-o", "--with-orbitals", action="store_true",
                        default=False)
    parser.add_argument("-c", "--with-cells", action="store_true",
                        default=False)
    parser.add_argument("-j", "--with-conj",  action="store_true",
                        default=False)
    parser.add_argument("-a", "--hop-as-arrows", action="store_true",
                        default=False)
    parser.add_argument("-e", "--hop-eng-cutoff", type=float, action="store",
                        default=1e-5)
    parser.add_argument("-v", "--verbose", action="store_true", default=False)
    parser.add_argument("model", type=str, action="store", default="model")
    args = parser.parse_args()

    # Parse lattice vectors
    with open(args.model, "r") as inf:
        content = inf.readlines()
    lat_vec = extract(content, "lattice")
    lat_vec = np.array([[float(v) for v in row] for row in lat_vec])

    # Parse lattice origin
    origin = extract(content, "origin")
    origin = np.array([float(v) for v in origin[0]])

    # Parse orbitals
    orb_info = extract(content, "orbitals")
    orb_pos = np.array([[float(v) for v in row[:3]] for row in orb_info])
    orb_eng = np.array([float(row[3]) for row in orb_info])
    orb_label = np.array([int(row[4]) for row in orb_info])

    # Parse hopping terms
    try:
        dr_info = extract(content, "dr")
    except ValueError:
        model_is_ample = False
        dr = None
        hop_info = extract(content, "hopping")
        hop_ind = np.array([[int(v) for v in row[:5]] for row in hop_info])
        hop_real = np.array([float(row[5]) for row in hop_info])
        hop_imag = np.array([float(row[6]) for row in hop_info])
        hop_eng = hop_real + 1j * hop_imag
    else:
        model_is_ample = True
        dr = np.array([[float(v) for v in row] for row in dr_info])
        hop_info = extract(content, "hopping")
        hop_ind = np.array([[int(v) for v in row[:2]] for row in hop_info])
        hop_real = np.array([float(row[2]) for row in hop_info])
        hop_imag = np.array([float(row[3]) for row in hop_info])
        hop_eng = hop_real + 1j * hop_imag

    # Build the model
    if model_is_ample:
        model = FakeCachedSample()
        model.sc0_lat_vec = lat_vec
        model.sc0_origin = origin
        model.orb_pos = orb_pos
        model.orb_eng = orb_eng
        model.hop_ind = hop_ind
        model.hop_eng = hop_eng
        model.dr = dr
        plot_class = tb.CachedSample
    else:
        model = tb.PrimitiveCell(lat_vec, origin, unit=tb.NM)
        for i in range(orb_pos.shape[0]):
            model.add_orbital(orb_pos[i], orb_eng.item(i), orb_label.item(i))
        for i in range(hop_ind.shape[0]):
            rn = hop_ind[i, :3]
            orb_i, orb_j = hop_ind.item(i, 3), hop_ind.item(i, 4)
            eng = hop_eng.item(i)
            model.add_hopping(rn, orb_i, orb_j, eng)
        plot_class = tb.PrimitiveCell

    plot_class.plot(model, with_orbitals=args.with_orbitals, with_cells=args.with_cells,
                    hop_eng_cutoff=args.hop_eng_cutoff, hop_as_arrows=args.hop_as_arrows)
    if not model_is_ample and args.verbose:
        model.print()


if __name__ == "__main__":
    main()
