import jax
import jax.numpy as jnp
import numpy as np
import pyscf
import chex
from jaxtyping import Float, Array, Int
from jsonargparse import CLI, Namespace
from functools import partial
from collections import namedtuple
from icecream import ic
from pyscf_ipu.nanoDFT import utils
from pyscf_ipu.exchange_correlation.b3lyp import b3lyp
from pyscf_ipu.electron_repulsion.direct  import (prepare_electron_repulsion_integrals, electron_repulsion_integrals, ipu_einsum)

HARTREE_TO_EV = 27.2114079527
EPSILON_B3LYP = 1e-20
HYB_B3LYP = 0.2

def energy(density_matrix, H_core, diff_JK, E_xc, E_nuc, _np=jax.numpy):
    """Density Functional Theory (DFT) solves the optimisation problem:

        min_{density_matrix} energy(density_matrix, ...)

    We like to think of `energy(...)` as a loss function. `density_matrix`
    represents the density of electrons as:

        rho(r) = sum_{ij}^N density_matrix_{ij} X_i(r) X_j(r)  where X_i(r)~exp(-|r|^2).

    Here N is the number of atomic orbitals (AO) **and** molecular orbitals (N=66 for C6H6).
    All input matrices (density_matrix, H_core, diff_JK) are (N, N). The X_i(r) are called
    Gaussian Type Orbitals (GTO). The inputs (diff_JK, E_xc) depend on density_matrix.
    """
    E_core = _np.sum(density_matrix * H_core)             # float = -712.04[Ha] for C6H6.
    E_J_K    = _np.sum(density_matrix * diff_JK)          # NOTE: diff_JK is already diff_JK = J - (K / 2 * HYB_B3LYP)
    E      = E_core + E_J_K/2 + E_xc + E_nuc              # float = -232.04[Ha] for C6H6.

    return _np.array([E, E_core, E_J_K/2, E_xc, E_nuc])   # Energy (and its terms).

def nanoDFT_iteration(i, vals, opts, mol):
    """Each call updates density_matrix attempting to minimize energy(density_matrix, ... ). """
    density_matrix, V_xc, diff_JK, O, H_core, L_inv                 = vals[:6]                  # All (N, N) matrices
    E_nuc, occupancy, ERI, grid_weights, grid_AO, diis_history, log = vals[6:]                  # Varying types/shapes.

    if opts.v:
        print("---------- MEMORY CONSUMPTION ----------")
        MB = 0
        for t in vals:
            try:
                if type(t) != type(()) and len(np.shape(t)) > 0:
                    print(t.nbytes/10**6, t.shape, t.dtype)
                    MB += t.nbytes/10**6
            except:
                print(type(t))
        print("ERI")
        for a in ERI: # prints weird in dense_ERI case
            print( a.nbytes/10**6, a.shape)
            MB += a.nbytes / 10**6
        print("__________")
        print("Total: ", MB)
        print("----------------------------------------")
        print("")

    # Step 1: Update Hamiltonian (optionally use DIIS to improve DFT convergence).
    H = H_core + diff_JK + V_xc                                                                 # (N, N)
    if opts.diis: H, diis_history = DIIS(i, H, density_matrix, O, diis_history, opts)           # H_{i+1}=c_1H_i+...+c9H_{i-9}.

    # Step 2: Solve eigh (L_inv turns generalized eigh into eigh).
    eigvects = L_inv.T @ linalg_eigh(L_inv @ H @ L_inv.T, opts)[1]                              # (N, N)

    # Step 3: Use result from eigenproblem to update density_matrix.
    density_matrix = (eigvects*occupancy*2) @ eigvects.T                                        # (N, N)
    E_xc, V_xc     = exchange_correlation(density_matrix, grid_AO, grid_weights)                # float (N, N)
    diff_JK        = get_JK(density_matrix, ERI, opts.dense_ERI, opts.backend)                  # (N, N) (N, N)

    # Log SCF matrices and energies (not used by DFT algorithm).
    #log["matrices"] = log["matrices"].at[i].set(jnp.stack((density_matrix, J, K, H)))           # (iterations, 4, N, N)
    N = density_matrix.shape[0]
    log["matrices"] = jax.lax.dynamic_update_slice(log["matrices"], density_matrix.reshape(1, 1, N, N), (i, 0, 0, 0))
    log["matrices"] = jax.lax.dynamic_update_slice(log["matrices"], diff_JK.       reshape(1, 1, N, N), (i, 1, 0, 0))
    log["matrices"] = jax.lax.dynamic_update_slice(log["matrices"], H.             reshape(1, 1, N, N), (i, 2, 0, 0))
    log["energy"]   = log["energy"].at[i].set(energy(density_matrix, H_core, diff_JK, E_xc, E_nuc))  # (iterations, 6)

    if opts.vis_num_error is True:
        import os
        dir_label = opts.molecule_name
        num_error_dir = f'num_error/{dir_label}/'
        os.makedirs(num_error_dir , exist_ok=True)

        def host_callback(data, i):
            # labels are adjusted to the `data` that will be passed to the callback - keep that in mind when passing different list of tensors
            labels = ["density_matrix", "V_xc", "diff_JK", "O", "H_core", "L_inv", "E_nuc", "occupancy", "ERI", "grid_weights", "grid_AO", "diis_history", "E_xc", "eigvects", "H"]
            for l, d  in zip(labels, data):
                if l == "diis_history" or l == "ERI":
                    for idx, arr in enumerate(d):
                        np.savez(f'{num_error_dir}{i}_{l}{idx}.npz', v = np.array(arr))
                else:
                    np.savez(f'{num_error_dir}{i}_{l}.npz', v = d)

        jax.debug.callback(host_callback, vals[:-1] + [E_xc, eigvects, H], i)

    return [density_matrix, V_xc, diff_JK, O, H_core, L_inv, E_nuc, occupancy, ERI, grid_weights, grid_AO, diis_history, log]

def exchange_correlation(density_matrix, grid_AO, grid_weights):
    """Compute exchange correlation integral using atomic orbitals (AO) evalauted on a grid. """
    # Perfectly SIMD parallelizable over grid_size axis.
    # Only need one reduce_sum in the end.
    grid_AO_dm = grid_AO[0] @ density_matrix                                                    # (gsize, N) @ (N, N) -> (gsize, N)
    grid_AO_dm = jnp.expand_dims(grid_AO_dm, axis=0)                                            # (1, gsize, N)
    mult = grid_AO_dm * grid_AO
    rho = jnp.sum(mult, axis=2)                                                                 # (4, grid_size)=(4, 45624) for C6H6.
    E_xc, vrho, vgamma = b3lyp(rho, EPSILON_B3LYP)                                              # (gridsize,) (gridsize,) (gridsize,)
    E_xc = jax.lax.psum(jnp.sum(rho[0] * grid_weights * E_xc), axis_name="p")                   # float=-27.968[Ha] for C6H6 at convergence.
    rho = jnp.concatenate([vrho.reshape(1, -1)/2, 4*vgamma*rho[1:4]], axis=0) * grid_weights    # (4, grid_size)=(4, 45624)
    grid_AO_T = grid_AO[0].T                                                                    # (N, gsize)
    rho = jnp.expand_dims(rho, axis=2)                                                          # (4, gsize, 1)
    grid_AO_rho = grid_AO * rho                                                                 # (4, gsize, N)
    sum_grid_AO_rho = jnp.sum(grid_AO_rho, axis=0)                                              # (gsize, N)
    V_xc = grid_AO_T @ sum_grid_AO_rho                                                          # (N, N)
    V_xc = jax.lax.psum(V_xc, axis_name="p")                                                    # (N, N)
    V_xc = V_xc + V_xc.T                                                                        # (N, N)
    return E_xc, V_xc                                                                           # (float) (N, N)

def get_JK(density_matrix, ERI, dense_ERI, backend):
    """Computes the (N, N) matrices J and K. Density matrix is (N, N) and ERI is (N, N, N, N).  """
    N = density_matrix.shape[0]

    if dense_ERI:
        J = jnp.einsum('ijkl,ji->kl', ERI, density_matrix)                                       # (N, N)
        K = jnp.einsum('ijkl,jk->il', ERI, density_matrix)                                       # (N, N)
        diff_JK = J - (K / 2 * HYB_B3LYP)
    else:
        from pyscf_ipu.nanoDFT.sparse_symmetric_ERI import sparse_symmetric_einsum
        diff_JK = sparse_symmetric_einsum(ERI[0], ERI[1], density_matrix, backend)

    diff_JK = diff_JK.reshape(N, N)

    return diff_JK

def _nanoDFT(state, ERI, grid_AO, grid_weights, opts, mol):
    # Utilize the IPUs MIMD parallism to compute the electron repulsion integrals (ERIs) in parallel.
    #if opts.backend == "ipu": state.ERI = electron_repulsion_integrals(state.input_floats, state.input_ints, mol, opts.threads_int, opts.intv)
    #else: pass # Compute on CPU.
    grid_AO = jnp.transpose(grid_AO, (1,0,2)) # (padded_gsize/16, 4, N) -> (4, pgsize, N)

    # Precompute the remaining tensors.
    E_xc, V_xc = exchange_correlation(state.density_matrix, grid_AO, grid_weights) # float (N, N)
    diff_JK    = get_JK(state.density_matrix, ERI, opts.dense_ERI, opts.backend)                      # (N, N) (N, N)
    H_core     = state.kinetic + state.nuclear                                           # (N, N)

    # Log matrices from all DFT iterations (not used by DFT algorithm).
    N = H_core.shape[0]
    log = {"matrices": np.zeros((opts.its, 4, N, N)), "E_xc": np.zeros((opts.its)), "energy": np.zeros((opts.its, 5))}

    # Perform DFT iterations.
    log = jax.lax.fori_loop(0, opts.its, partial(nanoDFT_iteration, opts=opts, mol=mol), [state.density_matrix, V_xc, diff_JK, state.O, H_core, state.L_inv,  # all (N, N) matrices
                                                            state.E_nuc, state.mask, ERI, grid_weights, grid_AO, state.diis_history, log])[-1]

    return log["matrices"], H_core, log["energy"]

FloatN = Float[Array, "N"]
FloatNxN = Float[Array, "N N"]
Grid = Float[Array, "4 grid_size N"]
FloatArray = Float[Array, "..."]
IntArray = Int[Array, "..."]

@chex.dataclass
class IterationState:
    """State tensors used during self-consistent field (SCF) iterations

        FloatN: Vector used to store the electron occupation mask.
        FloatNxN: Square matrix used for storing the one-electron integrals and density matrix.
        FloatNxNxNxN: 4-d matrix representing the two-electron repulsion integrals.
        Grid [4, grid_size, N] (float): Numerical grid used to evaluate the
            exchange-correlation energy integral.

    Attributes:
        E_nuc (float): Energy of the nuclear-nuclear electrostatic interactions.
        density_matrix (FloatNxN): Electron density in the LCAO basis set.
        kinetic (FloatNxN): Kinetic energy integrals in the LCAO basis set.
        nuclear (FloatNxN): nuclear attraction integrals in the LCAO basis set.
        O (FloatNxN): Overlap integrals in the LCAO basis set.
        mask (FloatN): Orbital occupation mask.
        input_floats (FloatArray): Supplementary vector of floats for ERI evaluation with libcint
        input_ints (IntArray): Supplementary vector of ints for ERI evaluation with libcint
        L_inv (FloatNxN): Defined as the inverse of the Cholesky decomposition of the overlap matrix.
            Used to change generalised eig problem into an eigh one.
        diis_history (FloatArray): Direct Inversion of Iterative Subspace (DIIS) is an optional method that
            can accelerate convergence of the SCF iterations. Maintains a history of how the Hamiltonian
            is evolving across the SCF iterations.

    """
    E_nuc: float
    density_matrix: FloatNxN
    kinetic: FloatNxN
    nuclear: FloatNxN
    O: FloatNxN
    mask: FloatN
    input_floats: FloatArray
    input_ints: IntArray
    L_inv: FloatNxN
    diis_history: FloatArray

def init_dft_tensors_cpu(mol, opts, DIIS_iters=9):
    N                = mol.nao_nr()                                 # N=66 for C6H6 (number of atomic **and** molecular orbitals)
    print("-----> [ %i ] <-----"%N)
    n_electrons_half = mol.nelectron//2                             # 21 for C6H6
    E_nuc            = mol.energy_nuc()                             # float = 202.4065 [Hartree] for C6H6. TODO(): Port to jax.

    # TODO(): port grid/eval_gto to Jax.
    grids            = pyscf.dft.gen_grid.Grids(mol)
    grids.level      = opts.level
    grids.build()
    grid_weights    = grids.weights                                 # (grid_size,) = (45624,) for C6H6
    coord_str       = 'GTOval_cart_deriv1' if mol.cart else 'GTOval_sph_deriv1'
    grid_AO         = mol.eval_gto(coord_str, grids.coords, 4)      # (4, grid_size, N) = (4, 45624, 9) for C6H6.
    if opts.ao_threshold > 0.0:
        grid_AO[np.abs(grid_AO)<opts.ao_threshold] = 0
        sparsity_mask = np.where(np.all(grid_AO == 0, axis=0), 0, 1)
        sparse_rows = np.where(np.all(sparsity_mask == 0, axis=1), 0, 1).reshape(-1, 1)
        print(f"axis=( ,  ) sparsity in grid_AO: {np.sum(grid_AO==0) / grid_AO.size:.4f}")
        print(f"axis=(0,  ) sparsity in grid_AO: {np.sum(sparsity_mask==0) / sparsity_mask.size:.4f}")
        print(f"axis=(0, 2) sparsity in grid_AO: {np.sum(sparse_rows==0) / sparse_rows.size:.4f}")
        grid_AO = jnp.delete(grid_AO, jnp.where(sparse_rows == 0)[0], axis=1)
        grid_weights = jnp.delete(grid_weights, jnp.where(sparse_rows == 0)[0], axis=0)
        grid_coords = jnp.delete(grids.coords, jnp.where(sparse_rows == 0)[0], axis=0)
    else:
        grid_coords = grids.coords
    density_matrix  = pyscf.scf.hf.init_guess_by_minao(mol)         # (N,N)=(66,66) for C6H6.

    # TODO(): Add integral math formulas for kinetic/nuclear/O/ERI.
    kinetic         = mol.intor_symmetric('int1e_kin')              # (N,N)
    nuclear         = mol.intor_symmetric('int1e_nuc')              # (N,N)
    O               = mol.intor_symmetric('int1e_ovlp')             # (N,N)
    L_inv           = np.linalg.inv(np.linalg.cholesky(O))          # (N,N)

    input_floats, input_ints = 0,0#prepare_electron_repulsion_integrals(mol)[:2]
    mask = np.concatenate([np.ones(n_electrons_half), np.zeros(N-n_electrons_half)])
    diis_history = None

    if opts.diis:
        # DIIS is an optional technique to improve DFT convergence.
        DIIS_H       = np.zeros((DIIS_iters+1, DIIS_iters+1))
        DIIS_H[0,1:] = DIIS_H[1:,0] = 1
        diis_history = (np.zeros((DIIS_iters, N**2)), np.zeros((DIIS_iters, N**2)), DIIS_H)

    state = IterationState(E_nuc=E_nuc, density_matrix=density_matrix, kinetic=kinetic,
                           nuclear=nuclear, O=O, mask=mask,
                           input_floats=input_floats, input_ints=input_ints,
                           L_inv=L_inv, diis_history=diis_history)

    return state, n_electrons_half, E_nuc, N, L_inv, grid_weights, grid_coords, grid_AO

def nanoDFT(mol, opts):
    # Init DFT tensors on CPU using PySCF.
    state, n_electrons_half, E_nuc, N, L_inv, _grid_weights, grid_coords, grid_AO = init_dft_tensors_cpu(mol, opts)

    grid_AO = jnp.transpose(grid_AO, (1, 0, 2)) # (4,gsize,N) -> (gsize,4,N)
    grid_weights = _grid_weights
    gsize = grid_AO.shape[0]

    remainder = gsize % opts.ndevices
    if remainder != 0:
        grid_AO = jnp.pad(grid_AO, ((0,remainder), (0,0), (0,0)) )
        grid_weights = jnp.pad(grid_weights, ((0,remainder)) )
    grid_AO = grid_AO.reshape(opts.ndevices, -1, 4, N)
    grid_weights = grid_weights.reshape(opts.ndevices, -1)

    # Run DFT algorithm (can be hardware accelerated).
    if opts.dense_ERI:
        assert opts.ndevices == 1, "Only support '--dense_ERI True' for `--ndevices 1`. "
        eri_in_axes = 0
        ERI = mol.intor("int2e_sph")
        ERI = np.expand_dims(ERI, 0)
        below_thr = np.abs(ERI) <= opts.eri_threshold
        ERI[below_thr] = 0.0
        ic(ERI.size, np.sum(below_thr), np.sum(below_thr)/ERI.size)
    else:
        from pyscf_ipu.nanoDFT.sparse_symmetric_ERI import get_i_j, num_repetitions_fast
        distinct_ERI         = mol.intor("int2e_sph", aosym="s8")
        print(distinct_ERI.size)
        below_thr = np.abs(distinct_ERI) <= opts.eri_threshold
        distinct_ERI[below_thr] = 0.0
        ic(distinct_ERI.size, np.sum(below_thr), np.sum(below_thr)/distinct_ERI.size)
        nonzero_indices      = np.nonzero(distinct_ERI)[0].astype(np.uint64)
        nonzero_distinct_ERI = distinct_ERI[nonzero_indices].astype(np.float32)

        ij, kl               = get_i_j(nonzero_indices)
        rep                  = num_repetitions_fast(ij, kl)
        nonzero_distinct_ERI = nonzero_distinct_ERI / rep
        batches  = int(opts.batches) # perhaps make 10 batches?
        nipu = opts.ndevices
        remainder = nonzero_indices.shape[0] % (nipu*batches)

        if remainder != 0:
            print(nipu*batches-remainder, ij.shape)
            ij = np.pad(ij, ((0,nipu*batches-remainder)))
            kl = np.pad(kl, ((0,nipu*batches-remainder)))
            nonzero_distinct_ERI = np.pad(nonzero_distinct_ERI, (0,nipu*batches-remainder))

        ij = ij.reshape(nipu, batches, -1)
        kl = kl.reshape(nipu, batches, -1)
        nonzero_distinct_ERI = nonzero_distinct_ERI.reshape(nipu, batches, -1)

        i, j = get_i_j(ij.reshape(-1))
        k, l = get_i_j(kl.reshape(-1))
        nonzero_indices = np.vstack([i,j,k,l]).T.reshape(nipu, batches, -1, 4).astype(np.int16)
        nonzero_indices = jax.lax.bitcast_convert_type(nonzero_indices, np.float16)

        ERI = [nonzero_distinct_ERI, nonzero_indices]
        eri_in_axes = [0,0]
    #jitted_nanoDFT = jax.jit(partial(_nanoDFT, opts=opts, mol=mol), backend=opts.backend)
    jitted_nanoDFT = jax.pmap(partial(_nanoDFT, opts=opts, mol=mol), backend=opts.backend,
                        in_axes=(None, eri_in_axes, 0, 0),
                        axis_name="p")
    print(grid_AO.shape, grid_weights.shape)
    vals = jitted_nanoDFT(state, ERI, grid_AO, grid_weights)
    logged_matrices, H_core, logged_energies = [np.asarray(a[0]).astype(np.float64) for a in vals] # Ensure CPU

    # It's cheap to compute energy/hlgap on CPU in float64 from the logged values/matrices.
    logged_E_xc = logged_energies[:, 3].copy()
    print(logged_energies[:, 0] * HARTREE_TO_EV)
    density_matrices, diff_JKs, H = [logged_matrices[:, i] for i in range(3)]
    energies, hlgaps = np.zeros((opts.its, 5)), np.zeros(opts.its)
    for i in range(opts.its):
        energies[i] = energy(density_matrices[i], H_core, diff_JKs[i], logged_E_xc[i], E_nuc, np)
        hlgaps[i]   = hlgap(L_inv, H[i], n_electrons_half, np)
    energies, logged_energies, hlgaps = [a * HARTREE_TO_EV for a in [energies, logged_energies, hlgaps]]
    mo_energy, mo_coeff = np.linalg.eigh(L_inv @ H[-1] @ L_inv.T)
    mo_coeff = L_inv.T @ mo_coeff
    return energies, (logged_energies, hlgaps, mo_energy, mo_coeff, grid_coords, _grid_weights)

def DIIS(i, H, density_matrix, O, diis_history, opts):
    # DIIS is an optional technique which improves DFT convergence by computing:
    #   H_{i+1} = c_1 H_i + ... + c_8 H_{i-8}  where  c=pinv(some_matrix)[0,:]
    # We thus like to think of DIIS as "fancy momentum".
    _V, _H, DIIS_H     = diis_history       # (diis_iters, N**2), (diis_iters, N**2), (diis_iters+1, diis_iters+1)
    diis_iters, d = _V.shape
    DIIS_head = i % _V.shape[0]             # int in {0, ..., diis_iters-1}
    sdf       = O @ density_matrix @ H      # (N, N)=(66,66) for C6H6.
    errvec    = sdf - sdf.T                 # (N, N)

    _V = jax.lax.dynamic_update_slice(_V, errvec.reshape(1, d), (DIIS_head, 0))     # (diis_iters, N**2)=(9, 4356) for C6H6.
    _H = jax.lax.dynamic_update_slice(_H, H.reshape(1, d),      (DIIS_head, 0))     # (diis_iters, N**2)

    mask = jnp.where(np.arange(_V.shape[0]) < jnp.minimum(i+1, _V.shape[0]), jnp.ones(_V.shape[0], dtype=_V.dtype), jnp.zeros(_V.shape[0], dtype=_V.dtype))
    tmps = (_V.reshape(diis_iters, 1, d) @ errvec.reshape(1, d, 1)).reshape(-1) * mask # (diis_iters, )

    DIIS_H = jax.lax.dynamic_update_slice( DIIS_H, tmps.reshape(1, -1), (DIIS_head+1, 1) ) # (diis_iters+1, diis_iters+1)
    DIIS_H = jax.lax.dynamic_update_slice( DIIS_H, tmps.reshape(-1, 1), (1, DIIS_head+1) ) # (diis_iters+1, diis_iters+1)

    mask_         = jnp.concatenate([jnp.ones(1, dtype=mask.dtype), mask]) # (diis_iters+1,)
    masked_DIIS_H = DIIS_H * mask_.reshape(-1, 1) * mask_.reshape(1, -1)

    if opts.backend == "ipu":  c = pinv0(masked_DIIS_H, opts)                 # (diis_iters+1,)=10 for C6H6.
    else:                      c = jnp.linalg.pinv(masked_DIIS_H)[0, :]       # (diis_iters+1,)=10 for C6H6.

    H = (c[1:] @ _H).reshape(H.shape)                                         # (N, N)
    return H, (_V, _H, DIIS_H)                                                # (N, N)

def hlgap(L_inv, H, n_electrons_half, _np):
    mo_energy   = _np.linalg.eigh(L_inv @ H @ L_inv.T)[0]
    return _np.abs(mo_energy[n_electrons_half] - mo_energy[n_electrons_half-1])

def linalg_eigh(x, opts):
    if opts.backend == "ipu" and x.shape[0] >= 6:
        from tessellate_ipu.linalg import ipu_eigh

        n = x.shape[0]
        pad = n % 2
        if pad:
            x = jnp.pad(x, [(0, 1), (0, 1)], mode='constant')
            #assert False

        eigvects, eigvals = ipu_eigh(x, sort_eigenvalues=True, num_iters=12)

        if pad:
            e1 = eigvects[-1:]
            col = jnp.argmax(e1)
            eigvects = jnp.roll(eigvects, -col-1)
            eigvects = eigvects[:, :-1]
            eigvects = jnp.roll(eigvects, -(-col))
            eigvects = eigvects[:-1]
            #assert False
    else:
        eigvals, eigvects = jnp.linalg.eigh(x)

    return eigvals, eigvects

def pinv0(a, opts):  # take out first row
    # TODO: add a comment explaining the role of this constant
    cond =  9*1.1920929e-07
    vals, vect = linalg_eigh(a, opts)
    c = vect @ ( jnp.where( jnp.abs(vals) > cond, 1/vals, 0) * vect[0, :])
    return c

def grad_elec(weight, grid_AO, eri, s1, h1aos, natm, aoslices, mask, mo_energy, mo_coeff, mol):
    # Electronic part of RHF/RKS gradients
    dm0  = 2 * (mo_coeff*mask) @ mo_coeff.T                                 # (N, N) = (66, 66) for C6H6.
    dme0 = 2 * (mo_coeff * mask*mo_energy) @  mo_coeff.T                    # (N, N) = (66, 66) for C6H6.

    # Code identical to exchange correlation.
    rho             = jnp.sum( grid_AO[:1] @ dm0 * grid_AO, axis=2)         # (10, grid_size) = (10, 45624) for C6H6.
    _, vrho, vgamma = b3lyp(rho, EPSILON_B3LYP)                             # (grid_size,) (grid_size,)
    V_xc            = jnp.concatenate([vrho.reshape(1, -1)/2, 4*vgamma.reshape(1, -1)*rho[1:4]], axis=0)  # (4, grid_size)

    vmat = grid_AO[1:4].transpose(0, 2, 1) @ jnp.sum(grid_AO[:4] * jnp.expand_dims(weight * V_xc, axis=2), axis=0) # (3, N, N)
    aos = jnp.concatenate([jnp.expand_dims(grid_AO[np.array([1,4,5,6])], 0), jnp.expand_dims(grid_AO[np.array([2,5,7,8])], 0), jnp.expand_dims(grid_AO[np.array([3,6,8,9])], 0)], axis=0) # (3, N, N)
    V_xc = - vmat - jnp.transpose(jnp.einsum("snpi,np->spi", aos, weight*V_xc), axes=(0,2,1)) @ grid_AO[0]  # (3, 4, grid_size, N)

    vj = - jnp.einsum('sijkl,lk->sij', eri, dm0) # (3, N, N)
    vk = - jnp.einsum('sijkl,jk->sil', eri, dm0) # (3, N, N)
    vhf = V_xc + vj - vk * .5 * HYB_B3LYP        # (3, N, N)

    de = jnp.einsum('lxij,ij->lx', h1aos, dm0)   # (natm, 3)
    for k, ia in enumerate(range(natm)):
        p0, p1 = aoslices[ia][2], aoslices[ia][3]
        de = de.at[k].add(jnp.einsum('xij,ij->x', vhf[:, p0:p1], dm0[p0:p1]) * 2)
        de = de.at[k].add(-jnp.einsum('xij,ij->x', s1[:, p0:p1], dme0[p0:p1]) * 2)
    return de

def grad_nuc(charges, coords):
    # Derivatives of nuclear repulsion energy wrt nuclear coordinates
    natm = charges.shape[0]
    pairwise_charges    = charges.reshape(natm, 1) * charges.reshape(1, natm)                # (natm, natm)
    pairwise_difference = coords.reshape(1, natm, 3) - coords.reshape(natm, 1, 3)            # (natm, natm, 3)
    pairwise_distances  = jnp.linalg.norm(pairwise_difference, axis=2) ** 3                  # (natm, natm)
    pairwise_distances  = jnp.where(pairwise_distances == 0, jnp.inf, pairwise_distances)    # (natm, natm)
    all = - pairwise_charges.reshape(natm, natm, 1) * pairwise_difference                    # (natm, natm, 3)
    all = all / pairwise_distances.reshape(natm, natm, 1)                                    # (natm, natm, 3)
    all = all.at[jnp.diag_indices(natm)].set(0)                                              # (natm, natm, 3)
    return jnp.sum(all, axis=0)                                                              # (natm, natm)

def grad(mol, coords, weight, mo_coeff, mo_energy):
    print(coords.shape, weight.shape)
    # Initialize DFT tensors on CPU using PySCF.
    ao = pyscf.dft.numint.NumInt().eval_ao(mol, coords, deriv=2)
    eri = mol.intor("int2e_ip1")
    s1  = - mol.intor('int1e_ipovlp', comp=3)
    kin = - mol.intor('int1e_ipkin',  comp=3)
    nuc = - mol.intor('int1e_ipnuc',  comp=3)

    mask = np.ones(mol.nao_nr())
    mask[mol.nelectron//2:] = 0

    aoslices = mol.aoslice_by_atom()
    h1 = kin + nuc
    def hcore_deriv(atm_id, aoslices, h1): # <\nabla|1/r|>
        _, _, p0, p1 = aoslices[atm_id]
        with mol.with_rinv_at_nucleus(atm_id):
            vrinv = mol.intor('int1e_iprinv', comp=3) #
            vrinv *= -mol.atom_charge(atm_id)
        vrinv[:,p0:p1] += h1[:,p0:p1]
        return vrinv + vrinv.transpose(0,2,1)
    N = h1.shape[1] # (3, N , N)
    h1aos = np.zeros((mol.natm, 3, N, N))
    for k, ia in enumerate(range(mol.natm)):
        p0, p1 = aoslices[ia,2:]
        h1aos[k] = hcore_deriv(ia, aoslices, h1)

    charges = np.zeros((mol.natm))
    coords = np.zeros((mol.natm,3))
    for j in range(mol.natm):
        charges[j] = mol.atom_charge(j)
        coords[j]= mol.atom_coord(j)

    #_grad_elec = jax.jit(grad_elec, static_argnames=["aoslices", "natm"], backend="cpu")
    _grad_elec = grad_elec
    _grad_nuc = jax.jit(grad_nuc, backend="cpu")

    return _grad_elec(weight, ao, eri, s1, h1aos, mol.natm, tuple([tuple(a) for a in aoslices.tolist()]), mask, mo_energy, mo_coeff, mol)  + _grad_nuc(charges, coords)

def pyscf_reference(mol_str, opts):
    from pyscf import __config__
    __config__.dft_rks_RKS_grids_level = opts.level

    mol = build_mol(mol_str, opts.basis)
    mol.max_cycle = opts.its
    mf = pyscf.scf.RKS(mol)
    mf.max_cycle = opts.its
    mf.xc = opts.xc
    mf.diis_space = 9
    if not opts.diis:  #
        mf.diis_space = 0
        mf.diis = False
    pyscf_energies = []
    pyscf_hlgaps = []
    lumo         = mol.nelectron//2
    homo         = lumo - 1
    def callback(envs):
        pyscf_energies.append(envs["e_tot"]*HARTREE_TO_EV)
        hl_gap_hartree = np.abs(envs["mo_energy"][homo] - envs["mo_energy"][lumo]) * HARTREE_TO_EV
        pyscf_hlgaps.append(hl_gap_hartree)
        print("\rPYSCF: ", pyscf_energies[-1] , end="")
    mf.callback = callback
    mf.kernel()
    print("")
    forces = mf.nuc_grad_method().kernel()
    return np.array(pyscf_energies), np.array(pyscf_hlgaps), np.array(forces)

def print_difference(nanoDFT_E, nanoDFT_forces, nanoDFT_logged_E, nanoDFT_hlgap, pyscf_E, pyscf_forces, pyscf_hlgap):
    #TODO(HH): rename to match caller variable names
    print("pyscf_hlgap\t%15f"%( pyscf_hlgap[-1]))
    print("us_hlgap\t%15f"%(    nanoDFT_hlgap[-1]))
    print("err_hlgap\t%15f"%np.abs(pyscf_hlgap[-1]  - nanoDFT_hlgap[-1]))
    print("pyscf:\t\t%15f"%pyscf_E[-1])
    print("us:\t\t%15f"%nanoDFT_E[-1, 0])
    print("mus:\t\t%15f"%np.mean(nanoDFT_E[-10:, 0]))
    print("diff:\t\t%15f"%np.abs(pyscf_E[-1]-nanoDFT_E[-1, 0]))
    print("mdiff:\t\t%15f"%np.abs(pyscf_E[-1]-np.mean(nanoDFT_E[-10:, 0])), np.std(nanoDFT_E[-10:, 0]))
    print("chemAcc: \t%15f"%0.043)
    print("chemAcc/diff: \t%15f"%(0.043/np.abs(pyscf_E[-1]-nanoDFT_E[-1, 0])))
    print("chemAcc/mdiff: \t%15f"%(0.043/np.abs(pyscf_E[-1]-np.mean(nanoDFT_E[-10:, 0]))))
    print("")
    pyscf_E = np.concatenate([pyscf_E, np.ones(nanoDFT_E.shape[0]-pyscf_E.shape[0])*pyscf_E[-1]])
    pyscf_hlgap = np.concatenate([pyscf_hlgap, np.ones(nanoDFT_hlgap.shape[0]-pyscf_hlgap.shape[0])*pyscf_hlgap[-1]])
    print("%18s"%"", "\t".join(["%10s"%str("iter %i "%i) for i in np.arange(1, nanoDFT_E.shape[0]+1)[1::3]]))
    print("%18s"%"Error Energy [eV]", "\t".join(["%10s"%str("%.2e"%f) for f in (pyscf_E[1::3] - nanoDFT_E[1::3, 0]).reshape(-1)]))
    print("%18s"%"Error HLGAP [eV]", "\t".join(["%10s"%str("%.2e"%f) for f in (pyscf_hlgap[1::3]   - nanoDFT_hlgap[1::3]).reshape(-1)]))

    print()
    print("%18s"%"E_core [eV]", "\t".join(["%10s"%str("%.5f"%f) for f in (nanoDFT_E[1::3, 1]).reshape(-1)]))
    print("%18s"%"E_J_K [eV]", "\t".join(["%10s"%str("%.5f"%f) for f in (nanoDFT_E[1::3, 2]).reshape(-1)]))
    print("%18s"%"E_xc [eV]", "\t".join(["%10s"%str("%.5f"%f) for f in (nanoDFT_E[1::3, 3]).reshape(-1)]))
    print("%18s"%"E_nuc [eV]", "\t".join(["%10s"%str("%.5f"%f) for f in (nanoDFT_E[1::3, 4]).reshape(-1)]))

    # Forces
    print()
    print("np.max(|nanoDFT_F-PySCF_F|):", np.max(np.abs(nanoDFT_forces-pyscf_forces)))

    norm_X = np.linalg.norm(nanoDFT_forces, axis=1)
    norm_Y = np.linalg.norm(pyscf_forces, axis=1)
    dot_products = np.sum(nanoDFT_forces * pyscf_forces, axis=1)
    cosine_similarity = dot_products / (norm_X * norm_Y)
    print("Force cosine similarity:",cosine_similarity)

def build_mol(mol_str, basis_name):
    mol = pyscf.gto.mole.Mole()
    mol.build(atom=mol_str, unit="Angstrom", basis=basis_name, spin=0, verbose=0)
    return mol

def nanoDFT_options(
        its: int = 20,
        mol_str: str = "benzene",
        float32: bool = False,
        basis: str = "sto-3g",
        xc: str = "b3lyp",
        backend: str = "cpu",
        level: int = 0,
        multv: int = 2,
        intv: int = 1,
        threads: int = 1,
        threads_int: int = 1,
        diis: bool = True,
        structure_optimization: bool = False, # AKA gradient descent on energy wrt nuclei
        eri_threshold : float = 0.0,
        ao_threshold: float = 0.0,
        batches: int = 32,
        ndevices: int = 1,
        dense_ERI: bool = False,
        v: bool = False, # verbose
        profile: bool = False, # if we only want profile exit after IPU finishes.
        vis_num_error: bool = False,
        molecule_name: str = None
):
    """
    nanoDFT

    Args:
        its (int): Number of Kohn-Sham iterations.
        mol_str (str): Molecule string, e.g., "H 0 0 0; H 0 0 1; O 1 0 0;" or one of:
            'benzene', 'methane', 'TRP', 'LYN', 'TYR', 'PHE', 'LEU', 'ILE', 'HIE', 'MET', 'GLN', 'HID', 'GLH', 'VAL', 'GLU', 'THR', 'PRO', 'ASN', 'ASH', 'ASP', 'SER', 'CYS',
            'CYX', 'ALA', 'GLY'
        float32 (bool) : Whether to use float32 (default is float64).
        basis (str): Which Gaussian basis set to use.
        xc (str): Exchange-correlation functional. Only support B3LYP
        backend (str): Accelerator backend to use: "--backend cpu" or "--backend ipu".
        level (int): Level of grids for XC numerical integration.
        gdb (int): Which version of GDP to load {10, 11, 13, 17}.
        multv (int): Which version of our einsum algorithm to use;comptues ERI@flat(v). Different versions trades-off for memory vs sequentiality
        intv (int): Which version to use of our integral algorithm.
        threads (int): For -backend ipu. Number of threads for einsum(ERI, dm) with custom C++ (trades-off speed vs memory).
        threads_int (int): For -backend ipu. Number of threads for computing ERI with custom C++ (trades off speed vs memory).
        eri_threshold (float): Zero out ERIs that are below the threshold in absolute value. Not supported for '--backend ipu'.
        ao_threshold (float): Zero out grid_AO that are below the threshold in absolute value.
        dense_ERI (bool): Whether to use dense ERI (s1) or sparse symmtric ERI.
    """
    if molecule_name is None:
        # use mol_str as a molecule name (in case it has not been provided)
        # before mol_str CLI arg is preprocessed and overwritten
        molecule_name = mol_str

    # From a compound name or CID, get a list of its atoms and their coordinates
    mol_str = utils.process_mol_str(mol_str)
    if mol_str is None:
        exit(1)

    print(f"Minimum interatomic distance: {utils.min_interatomic_distance(mol_str)}") # TODO: dies for --mol_str methane

    args = locals()
    mol_str = args["mol_str"]
    del args["mol_str"]
    args = Namespace(**args)

    from pyscf_ipu.experimental.device import has_ipu
    import os
    if has_ipu() and "JAX_IPU_USE_MODEL" in os.environ:
        args.dense_ERI = True
    args = namedtuple('DFTOptionsImmutable',vars(args).keys())(**vars(args)) # make immutable
    if not args.float32:
        jax.config.update('jax_enable_x64', not float32)
    return args, mol_str

def main():
    # Limit PySCF threads to mitigate problem with NUMA nodes.
    jax.config.FLAGS.jax_platform_name = 'cpu'
    import os
    opts, mol_str = CLI(nanoDFT_options)
    assert opts.xc == "b3lyp"
    print("Precision: float32") if opts.float32 else print("Precision: float64")

    if not opts.structure_optimization:
        # Test Case: Compare nanoDFT against PySCF.
        mol = build_mol(mol_str, opts.basis)

        print(f"Number of Atomic Orbitals\t{mol.nao_nr():15d}")
        print(f"Number of electrons\t{mol.nelectron:15d}")

        nanoDFT_E, (nanoDFT_logged_E, nanoDFT_hlgap, mo_energy, mo_coeff, grid_coords, grid_weights) = nanoDFT(mol, opts)
        nanoDFT_forces = grad(mol, grid_coords, grid_weights, mo_coeff, mo_energy)
        pyscf_E, pyscf_hlgap, pyscf_forces = pyscf_reference(mol_str, opts)
        print_difference(nanoDFT_E, nanoDFT_forces, nanoDFT_logged_E, nanoDFT_hlgap, pyscf_E, pyscf_forces, pyscf_hlgap)

        if opts.vis_num_error is True:
            from utils import save_plot
            import sys
            _plot_title = f"Created with:  python {' '.join(sys.argv)}"
            save_plot("num_error/", opts.molecule_name, opts.its, _plot_title)
    else:
        # pip install mogli imageio[ffmpeg] matplotlib
        import mogli
        import imageio
        import matplotlib.pyplot as plt
        opts = opts._replace(basis="6-31G")
        p = np.array([[0,1,1], [0,2,2], [0,3,3],
                      [0,4,4], [0,5,5], [0,6,6]])
        np.random.seed(42)
        p = p + np.random.normal(0, 0.3, p.shape) # slightly break symmetry
        A = ["H", "O", "H", "H", "O", "H"]
        natm = p.shape[0]
        os.makedirs("_tmp/", exist_ok=True)
        E = []
        ims = []
        for i in range(20):
            mol_str = "".join([f"{A[i]} {p[i]};".replace("[", "]").replace("]", "") for i in range(natm)])
            mol = build_mol(mol_str, opts.basis)
            nanoDFT_E, (nanoDFT_logged_E, nanoDFT_hlgap, mo_energy, mo_coeff, grid_coords, grid_weights) = nanoDFT(mol, opts)
            f = open(f"_tmp/{i}.xyz", "w")
            f.write(f"""{natm}\n{mol_str} {nanoDFT_E[-1, 0]}\n"""+"".join([f"{A[i]} {p[i]}\n".replace("[", "").replace("]", "") for i in range(natm)]))
            f.close()
            molecules = mogli.read(f'_tmp/{i}.xyz')
            mogli.export(molecules[0], f'_tmp/{i}.png', width=400, height=400,
                 bonds_param=1.15, camera=((9, 0, 0),
                                           (0, 0, 0),
                                           (0, 9, 0)))
            ims.append(imageio.v2.imread(f"_tmp/{i}.png/"))
            E.append(nanoDFT_E[-1, 0])
            nanoDFT_forces = grad(mol, grid_coords, grid_weights, mo_coeff, mo_energy)
            p = p - nanoDFT_forces
            print(nanoDFT_E[-1, 0], i)
        writer = imageio.get_writer('_tmp/test.gif', loop=0, duration=3)
        fig, ax = plt.subplots(1, 2, figsize=(10, 4))
        for c, i in enumerate(ims):
            for a in ax: a.cla()
            ax[0].axis("off")
            ax[1].set_ylabel("Energy [eV]")
            ax[1].set_xlabel("Step Number in Structure Optimization")
            ax[0].imshow(i)
            ax[1].plot(E, label="energy [eV]")
            ax[1].legend()
            ax[1].plot([c, c], [np.min(E), np.max(E)], '-k')
            plt.tight_layout()
            plt.savefig("_tmp/tmp.jpg")
            writer.append_data(imageio.v2.imread("_tmp/tmp.jpg"))
        writer.close()

if __name__ == "__main__":
    main()


"""
"""
import lvgl as lv
import ili9XXX
from ili9XXX import st7789
import machine
import time

# see pin_defs.py and import the pin defs that match your build
from pin_defs import dev_board as pins
# from pin_defs import manual_wiring as pins
import fs_driver

lv.init()

# FS driver init
fs_drv = lv.fs_drv_t()
fs_driver.fs_register(fs_drv, 'S')

# Load the font
opensans_regular_17 = lv.font_load("S:/opensans_regular_17.bin")

disp = st7789(
    **pins["st7789"],
    width=240, height=240, rot=ili9XXX.LANDSCAPE
)

scr = lv.scr_act()
scr.clean()

scr.set_style_bg_color(lv.color_hex(0x000000), 0)
scr.set_scrollbar_mode(lv.SCROLLBAR_MODE.OFF)

EDGE_PADDING = 8
COMPONENT_PADDING = 8

### Top Nav ###
top_nav = lv.obj(scr)
top_nav.set_size(240, 48)
top_nav.align(lv.ALIGN.TOP_LEFT, 0, 0)

style = lv.style_t()
style.init()
style.set_pad_all(0)
style.set_border_width(0)
style.set_radius(0)
style.set_bg_color(lv.color_hex(0x0000ff))
top_nav.add_style(style, 0)

label = lv.label(top_nav)
label.set_text("Settings")
label.center()
label_style = lv.style_t()
label_style.set_text_font(opensans_regular_17)
label_style.set_text_color(lv.color_hex(0xffffff))
label.add_style(label_style, 0)

obj = lv.btn(scr)
obj.set_size(240 - 2*EDGE_PADDING, 32)
obj.add_flag(lv.obj.FLAG.SNAPPABLE)
obj.center()

style = lv.style_t()
style.set_bg_color(lv.color_hex(0xffa500))
style.set_text_color(lv.color_hex(0x000000))
style.set_pad_all(0)
obj.add_style(style, 0)

label = lv.label(obj)
label.set_text("")
label.center()

key1 = machine.Pin(pins["buttons"]["key1"], machine.Pin.IN, machine.Pin.PULL_UP)
key2 = machine.Pin(pins["buttons"]["key2"], machine.Pin.IN, machine.Pin.PULL_UP)
key3 = machine.Pin(pins["buttons"]["key3"], machine.Pin.IN, machine.Pin.PULL_UP)

joy_up = machine.Pin(pins["buttons"]["joy_up"], machine.Pin.IN, machine.Pin.PULL_UP)
joy_down = machine.Pin(pins["buttons"]["joy_down"], machine.Pin.IN, machine.Pin.PULL_UP)
joy_left = machine.Pin(pins["buttons"]["joy_left"], machine.Pin.IN, machine.Pin.PULL_UP)
joy_right = machine.Pin(pins["buttons"]["joy_right"], machine.Pin.IN, machine.Pin.PULL_UP)
joy_press = machine.Pin(pins["buttons"]["joy_press"], machine.Pin.IN, machine.Pin.PULL_UP)

buttons = [
    (key1, "KEY1"),
    (key2, "KEY2"),
    (key3, "KEY3"),
    (joy_up, "UP"),
    (joy_down, "DOWN"),
    (joy_left, "LEFT"),
    (joy_right, "RIGHT"),
    (joy_press, "PRESS"),
]

cur_button = None

while True:
    for button, btn_name in buttons:
        if cur_button != button:
            if button.value() == 0:
                cur_button = button
                label.set_text(btn_name)
                break
        if cur_button == button and button.value() == 1:
            cur_button = None
    time.sleep(0.01)


from game.Params import *
from game.BonusManager import Bonus

BONUS_IMAGES = {Bonus.Pause: {YELLOW: 'game/images/pause_yellow.png',
                              GREEN: 'game/images/pause_green.png',
                              BLUE: 'game/images/pause_blue.png',
                              RED: 'game/images/pause_red.png'},
                Bonus.Reverse: {YELLOW: 'game/images/reverse_yellow.png',
                                GREEN: 'game/images/reverse_green.png',
                                BLUE: 'game/images/reverse_blue.png',
                                RED: 'game/images/reverse_red.png'},
                Bonus.Bomb: {YELLOW: 'game/images/bomb_yellow.png',
                             GREEN: 'game/images/bomb_green.png',
                             BLUE: 'game/images/bomb_blue.png',
                             RED: 'game/images/bomb_red.png'},
                Bonus.Speed: {YELLOW: 'game/images/speed_yellow.png',
                              GREEN: 'game/images/speed_green.png',
                              BLUE: 'game/images/speed_blue.png',
                              RED: 'game/images/speed_red.png'}}

class Button:
    def __init__(self, button_title, position, width=BTN_WIDTH,
                 height=BTN_HEIGHT, background_color=BROWN, font_color=TAUPE):
        self.title = button_title
        self.font = pygame.font.Font('game/fonts/Azov.ttf', FONT_SIZE)
        self.title_width, self.title_height = self.font.size(self.title)
        self.center = (position[0], position[1])
        self.width, self.height = width, height
        self.x_start, self.y_start = self.center[0] - self.width // 2, \
                                     self.center[1] - self.height // 2
        self.rect = pygame.Rect((self.x_start, self.y_start,
                                 width, height))
        self.background_color = background_color
        self.font_color = font_color

class Label:
    def __init__(self, text, position, color=BROWN):
        self.font = pygame.font.Font('game/fonts/Azov.ttf', FONT_SIZE)
        self.color = color
        self.text = self.font.render(text, True, color)
        self.width, self.height = self.font.size(text)
        self.x_start, self.y_start = position[0] - self.width // 2, \
                                     position[1] - self.height // 2

class Display:
    def __init__(self, background_color=TAUPE, buttons=None, labels=None,
                 sprites=None):
        if buttons is None:
            self.buttons = []
        else:
            self.buttons = buttons

        if sprites is None:
            self.spites = []
        else:
            self.spites = sprites

        if labels is None:
            self.labels = []
        else:
            self.labels = labels

        self.background_color = background_color

class UiManager:
    def __init__(self, screen, level):
        self.screen = screen
        self.level = level

        self.start_game_btn = Button('Начать игру', SCREEN_CENTER)
        self.start_game_display = Display(buttons=[self.start_game_btn])

        self.level_label = Label('Уровень {}'.format(level.number),
                                   (WIDTH // 2, 40))
        sprites = [level.player, level.path, level.ball_generator,
                   level.finish, level.shooting_manager]
        self.game_display = Display(sprites=[sprite for sprite in sprites],
                                    labels=[self.level_label])

        self.continue_btn = Button('Продолжить', SCREEN_CENTER)
        self.win_level_display = Display(buttons=[self.continue_btn])

        self.start_level_again_btn = Button('Начать сначала', SCREEN_CENTER,
                                            background_color=TAUPE,
                                            font_color=BROWN)
        self.lose_level_display = Display(BROWN,
                                          buttons=[self.start_level_again_btn])

        self.finish_btn = Button('Закончить', (WIDTH // 2, HEIGHT // 2 +
                                               2 * BTN_HEIGHT))
        self.start_game_again_btn = Button('Начать сначала', SCREEN_CENTER)
        self.win_label = Label('Вы прошли игру!', (WIDTH // 2, HEIGHT // 2 -
                                               2 * BTN_HEIGHT))
        self.win_game_display = Display(buttons=[self.start_game_again_btn,
                                                 self.finish_btn],
                                        labels=[self.win_label])

        self.new_game_button = Button('Новая игра', SCREEN_CENTER,
                                      background_color=TAUPE,
                                      font_color=BROWN)
        self.lose_game_display = Display(BROWN,
                                         buttons=[self.new_game_button])

    def draw_button(self, button):
        width, height = button.width, button.height
        x_start, y_start = button.x_start, button.y_start
        title_params = (x_start + width / 2 - button.title_width / 2,
                        y_start + height / 2 - button.title_height / 2)
        pygame.draw.rect(self.screen, button.background_color,
                         (x_start, y_start, width, height))
        self.screen.blit(button.font.render(button.title, True,
                                            button.font_color), title_params)
        button.rect = pygame.Rect((x_start, y_start, width, height))

    def draw_window(self, window):
        self.screen.fill(window.background_color)
        for button in window.buttons:
            self.draw_button(button)
        for label in window.labels:
            self.put_label(label)
        for sprite in window.spites:
            sprite.draw(self.screen)

    def show_score(self, points):
        points_label = Label('Очки: {}'.format(points), (WIDTH // 4, 40))
        self.put_label(points_label)

    def show_lives(self, lives):

        self.put_label(Label(str(lives), (3 * WIDTH // 4 + 30, 40)))
        self.screen.blit(pygame.transform.smoothscale(
            pygame.image.load("game/images/life.png"), (20, 20)),
            (3 * WIDTH // 4, 30))

    def put_label(self, label, color=TAUPE):
        pygame.draw.rect(self.screen, color, (label.x_start - label.width / 2,
                                              label.y_start, label.width,
                                              label.height))
        self.screen.blit(label.text, (label.x_start, label.y_start))


import csv
import math
import os
import sys
import torch
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
import Feedforward
import Kalman
from BKF_testing import start_test
from CircleGenerator import CirclesDataset, create_directory, save_tensor
from DeviceDataLoader import DeviceDataLoader, get_default_device
from datetime import datetime

device = get_default_device()

def get_val_loss(FFModel, KalmanModel, val_loader, loss_function, random_start_pos, random_start_vel, upper_bound_L_hat,
                 simplified_cov_update):
    with torch.no_grad():
        FFModel.eval()
        KalmanModel.eval()
        total_loss = 0.0
        num_batches = 0
        for batch in val_loader:
            num_batches += 1
            loss, _, _ = end_to_end_training_step(FFModel, KalmanModel, batch, loss_function, random_start_pos,
                                                  random_start_vel, upper_bound_L_hat, simplified_cov_update)
            total_loss += loss
        return total_loss / num_batches

def epoch_end(self, epoch, result):
    print("Epoch [{}], train_loss: {:.4f}, val_loss: {:.4f}, last_lr: {:.4f}".format(
        epoch, result['train_loss'], result['val_loss'], result['lr'][-1]))

def epoch_end_no_lr(self, epoch, result):
    print("Epoch [{}], train_loss: {:.4f}, val_loss: {:.4f}".format(
        epoch, result['train_loss'], result['val_loss']))

def apply_upper_bound(L_hat_List):
    # L_hat_List tensor of shape (N * T, 3)
    upper_bound = 2.0
    lower_bound = -2.0
    reference_tensor_up = torch.tensor([upper_bound, math.exp(upper_bound), upper_bound]) \
        .repeat(L_hat_List.shape[0], 1).to(device)

    reference_tensor_down = torch.tensor([lower_bound, -math.exp(lower_bound), lower_bound]) \
        .repeat(L_hat_List.shape[0], 1).to(device)

    condition = L_hat_List > reference_tensor_up
    L_hat_List[condition] = upper_bound

    condition = L_hat_List < reference_tensor_down
    L_hat_List[condition] = lower_bound

def end_to_end_training_step(FFModel, KalmanModel, sample_batched, loss_function, random_start_pos, random_start_vel,
                             upper_bound_L_hat, simplified_cov_update):
    # N is batch size
    # T is amt of frames in sequence
    # Sample dimensions (images, positions), images: (N, T, channels, width, height), labels: (N, T, (x, y, vx, vy) )
    images = sample_batched[0].permute(1, 0, 2, 3, 4)
    positions = sample_batched[1].permute(1, 0, 2)[:, :, :2]
    vels = sample_batched[1].permute(1, 0, 2)[:, :, 2:]

    # Format data
    # Images in sequence, not including first image (T-1, N, 3, 128, 128)
    images = torch.stack([images[ii] for ii in range(1, len(images))]).float().to(device)

    first_frame_positions = None
    first_frame_vels = None
    if not random_start_pos:
        # First frames of every sequence in batch (N, 2)
        first_frame_positions = positions[0]
    if not random_start_vel:
        # First frames of every sequence in batch (N, 2)
        first_frame_vels = vels[0]

    # (T, N, 2)
    positions = torch.stack([positions[ii] for ii in range(1, len(positions))]).float().to(device)

    # Reshape images so everything can be processed in parallel by utilizing batch size
    T, N, C, W, H = images.shape[0], images.shape[1], images.shape[2], images.shape[3], images.shape[4]
    seq_images = images.view(T * N, C, W, H)

    # Forward pass
    # output (T * N, dim_output)
    z_list, L_hat_list = FFModel(seq_images)

    if upper_bound_L_hat:
        apply_upper_bound(L_hat_list)

    # Decompress the results into original images format
    z_list = z_list.view(T, N, z_list.shape[1])
    L_hat_list = L_hat_list.view(T, N, L_hat_list.shape[1])

    # Pass through KF
    position_prediction = KalmanModel(z_list, L_hat_list, first_frame_positions, first_frame_vels,
                                      simplified_cov_update)

    loss = loss_function(position_prediction, positions)
    return loss, position_prediction, L_hat_list

def fit(epochs, lr, FFModel, KalmanModel, train_loader, val_loader, opt_func, loss_function, on_server, save_matrices,
        save_predictions, save_lrs, lr_scheduling, random_start_pos, random_start_vel, path_model,
        upper_bound_L_hat, simplified_cov_update):
    history = []
    predictionsDict = dict()
    matricesDict = dict()
    lrsDict = dict()

    optimizer = opt_func(list(FFModel.parameters()) + list(KalmanModel.parameters()), lr)

    sched = None
    lowest_val_loss = None
    loss = None
    model_dict = None

    if lr_scheduling:
        sched = torch.optim.lr_scheduler.OneCycleLR(optimizer, lr, epochs=epochs, steps_per_epoch=len(train_loader))

    for epoch in range(epochs):
        # Training Phase
        FFModel.train()
        KalmanModel.train()
        train_losses = []
        lrs = []
        for i, batch in enumerate(train_loader):
            loss, predictions, L_hat_list = end_to_end_training_step(FFModel, KalmanModel, batch, loss_function,
                                                                     random_start_pos,
                                                                     random_start_vel, upper_bound_L_hat,
                                                                     simplified_cov_update)

            train_losses.append(loss)
            loss.backward()
            optimizer.step()

            if lr_scheduling:
                lrs.append(get_lr(optimizer))
                sched.step()

            if i == int(len(train_loader.dl.dataset) / train_loader.dl.batch_size) - 1:
                # Save last tensor of last batch with its label and prediction
                if save_predictions:
                    tensor = batch[0][-1][-1]
                    label = batch[1][-1][-1]
                    prediction = predictions[-1][-1]
                    if "tensor" not in predictionsDict:
                        predictionsDict["tensor"] = tensor
                    predictionsDict[epoch] = (label, prediction)

                if save_matrices:
                    matricesDict[epoch] = ((KalmanModel.diag_Q.clone(),
                                            KalmanModel.diag_Q.grad.clone()), L_hat_list)

                if save_lrs:
                    lrsDict[epoch] = get_lr(optimizer)

                optimizer.zero_grad()

        # Validation phase
        mean_val_loss = get_val_loss(FFModel, KalmanModel, val_loader, loss_function, random_start_pos,
                                     random_start_vel, upper_bound_L_hat, simplified_cov_update)

        if lr_scheduling:
            result = {'train_loss': torch.stack(train_losses).mean().item(), 'val_loss': mean_val_loss, 'lr': lrs}
            epoch_end(epoch, result)
        else:
            result = {'train_loss': torch.stack(train_losses).mean().item(), 'val_loss': mean_val_loss}
            epoch_end_no_lr(epoch, result)

        history.append(result)

        if path_model is not None and epoch > epochs // 2:
            # Save model if lowest val loss
            val_loss = result['val_loss']
            if lowest_val_loss is None or lowest_val_loss > val_loss:
                lowest_val_loss = val_loss
                model_dict = {
                    "epoch": epoch,
                    "ff_state_dict": FFModel.state_dict(),
                    "kalman_state_dict": KalmanModel.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "loss": loss.item()
                }

    path = ""
    if on_server:
        path += f"/export/home2/NoCsBack/thesisdt/"
    done = datetime.now().strftime("%d-%m-%Y_%H-%M")
    path += f"WV/logs/bkf_logs_paper/t_{len(train_loader.dl.dataset)}_v_{len(val_loader.dl.dataset)}_b_{train_loader.dl.batch_size}_" \
            f"{done}/"

    if save_predictions or save_matrices:
        create_directory(path)

    if save_matrices:
        save_matrices_as_csv(matricesDict, path)

    if save_predictions:
        save_tensor_label_prediction(predictionsDict, path)

    if save_lrs:
        save_lrs_as_csv(lrsDict, path)

    if path_model is not None:
        done = datetime.now().strftime("%d-%m-%Y_%H-%M")
        model_dir = f"{path_model}_loss_{lowest_val_loss}_{done}"
        create_directory(model_dir)

        torch.save(model_dict, f"{model_dir}/lowest_val_loss.tar")

    return history, model_dir + "/lowest_val_loss.tar"

def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

def plot_losses(history):
    train_losses = [x.get('train_loss') for x in history]
    val_losses = [x.get('val_loss').item() for x in history]
    plt.figure()
    plt.plot(train_losses, '-bx')
    plt.plot(val_losses, '-rx')
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.legend(['Training', 'Validation'])
    plt.title('Loss vs. No. of epochs')
    plt.show()

def save_matrices_as_csv(matricesDict, path):
    csv_Q = "matrix_Q.csv"
    csv_L_hat = "matrix_L_hat"
    save_matrix_and_gradient_as_csv("Q", path, csv_Q, matricesDict, 0)
    save_matrix_as_csv("L_hat", path, csv_L_hat, matricesDict, 1)

def save_matrix_as_csv(name, path, csv_name, matricesDict, indexDict):
    with open(path + csv_name, "w") as csvfile:
        writer = csv.writer(csvfile)
        # Header: epoch | matrix name | gradient
        writer.writerow(["epoch", "matrix " + name])
        for epoch in range(0, len(matricesDict)):
            writer.writerow([epoch, matricesDict[epoch][indexDict].cpu().detach().numpy()])

def save_matrix_and_gradient_as_csv(name, path, csv_name, matricesDict, indexDict):
    with open(path + csv_name, "w") as csvfile:
        writer = csv.writer(csvfile)
        # Header: epoch | matrix name | gradient
        writer.writerow(["epoch", "matrix " + name, "gradient"])
        for epoch in range(0, len(matricesDict)):
            writer.writerow([epoch, matricesDict[epoch][indexDict][0].cpu().detach().numpy(),
                             matricesDict[epoch][indexDict][1].cpu().detach().numpy()])

def save_lrs_as_csv(lrsDict, path):
    csv_name = "lrs.csv"
    with open(path + csv_name, "w") as csvfile:
        writer = csv.writer(csvfile)
        # Header: epoch | lr
        writer.writerow(["epoch", "lr "])
        for epoch in range(0, len(lrsDict)):
            writer.writerow([epoch, lrsDict[epoch]])

def save_tensor_label_prediction(predictionDict, path):
    save_tensor(predictionDict["tensor"].cpu(), path, "tensor.pt")
    csv_name = "tensor_label_pred.csv"
    with open(path + csv_name, "w") as csvfile:
        writer = csv.writer(csvfile)
        # Header: epoch | label x | label y | prediction x | prediction y
        writer.writerow(["epoch", "label x", "label y", "prediction x", "prediction y"])
        # Skip the tensor entry
        for epoch in range(0, len(predictionDict) - 1):
            writer.writerow([epoch, predictionDict[epoch][0][0].item(), predictionDict[epoch][0][1].item(),
                             predictionDict[epoch][1][0].item(), predictionDict[epoch][1][1].item()])

def train_and_validate(nr_of_circles, nr_of_sequences_train, nr_of_sequences_val, nr_of_frames_train, nr_of_frames_val,
                       nr_of_epochs, batch_size, on_server, save_matrices, save_predictions, save_lrs, lr,
                       lr_scheduling,
                       random_start_pos, random_start_vel, static_seq, save_model, upper_bound_L_hat,
                       simplified_cov_update, change_shape, ff_model_path):
    train_csv_file = ''
    val_csv_file = ''
    if on_server:
        train_csv_file += '/export/home2/NoCsBack/thesisdt/WV/'
        val_csv_file += '/export/home2/NoCsBack/thesisdt/WV/'
    if static_seq:
        train_csv_file += f'datasets/seq_nonmoving_circles/'
        val_csv_file += f'datasets/seq_nonmoving_circles/'
    else:
        train_csv_file += f'datasets/linear_moving_circles_bkf/'
        val_csv_file += f'datasets/linear_moving_circles_bkf/'
    train_csv_file += f'circles={nr_of_circles}_frames={nr_of_frames_train}_noise=None' \
                      f'/train={nr_of_sequences_train}/train.csv'
    val_csv_file += f'circles={nr_of_circles}_frames={nr_of_frames_val}_noise=None' \
                    f'/validation={nr_of_sequences_val}/validation.csv'

    print("Training on: ")
    print(train_csv_file)
    print("----------------------------------")

    train_ds = CirclesDataset(train_csv_file)
    val_ds = CirclesDataset(val_csv_file)

    train_dl = DataLoader(train_ds, batch_size, pin_memory=True)
    val_dl = DataLoader(val_ds, batch_size, pin_memory=True)

    train_dl = DeviceDataLoader(train_dl, device)
    val_dl = DeviceDataLoader(val_dl, device)

    opt_func = torch.optim.Adam
    loss_func = torch.nn.MSELoss()

    FFModel = Feedforward.FFNetwork().to(device)
    KalmanModel = Kalman.KalmanFilter(device, batch_size).to(device)

    path_save_model = None
    if save_model:
        path_save_model = ""
        if on_server:
            path_save_model += f"/export/home2/NoCsBack/thesisdt/WV/"
        changing = ''
        if change_shape:
            changing = 'changing_'
        path_save_model += f"logs/models/bkf/{changing}circles={nr_of_circles}/t_{len(train_dl.dl.dataset)}_v_" \
                           f"{len(val_dl.dl.dataset)}_b_{batch_size}"

    # Load trained feedforward model
    if ff_model_path is not None:
        checkpoint = torch.load(ff_model_path)
        FFModel.load_state_dict(checkpoint['ff_state_dict'])
        print("Using trained FFModel")
        print(checkpoint['path_training_set'])
        print("----------------------------------")

    print("ff model", ff_model_path)
    print("nr of circles: ", nr_of_circles)
    print("nr of sequences train: ", nr_of_sequences_train)
    print("nr of sequences val: ", nr_of_sequences_val)
    print("lr: ", lr)
    print("lr sched: ", lr_scheduling)
    print("batch size: ", batch_size)
    print("epochs: ", nr_of_epochs)
    print("random_start_pos: ", random_start_pos)
    print("random_start_vel: ", random_start_vel)
    print("change shape: ", change_shape)

    history, model_path = fit(nr_of_epochs, lr, FFModel, KalmanModel, train_dl, val_dl, opt_func, loss_func, on_server,
                              save_matrices, save_predictions, save_lrs, lr_scheduling, random_start_pos,
                              random_start_vel,
                              path_save_model, upper_bound_L_hat, simplified_cov_update)
    plot_losses(history)

    return model_path

if __name__ == "__main__":
    on_server = False
    save_matrices = False
    save_predictions = False
    save_lrs = False
    random_start_pos = False
    random_start_vel = True
    static_seq = False
    save_model = True
    upper_bound_L_hat = False
    simplified_cov_update = False

    if len(sys.argv) == 12:
        nr_of_sequences_train, sequence_length_train, nr_of_sequences_val, sequence_length_val, nr_of_circles, \
        change_shape, nr_of_epochs, batch_size, ff_model_path, lr, lr_sched = sys.argv[1:]
        train_and_validate(nr_of_circles=int(nr_of_circles), nr_of_sequences_train=int(nr_of_sequences_train),
                           nr_of_sequences_val=int(nr_of_sequences_val), nr_of_frames_train=int(sequence_length_train),
                           nr_of_frames_val=int(sequence_length_val), nr_of_epochs=int(nr_of_epochs),
                           batch_size=int(batch_size), on_server=False, save_matrices=save_matrices,
                           save_predictions=save_predictions, save_lrs=save_lrs,
                           change_shape=(change_shape == 'change_shape'), lr=float(lr),
                           lr_scheduling=(lr_sched == 'lr_sched'), random_start_pos=random_start_pos,
                           random_start_vel=random_start_vel, static_seq=static_seq, save_model=True,
                           upper_bound_L_hat=False, simplified_cov_update=False, ff_model_path=ff_model_path)

    else:
        raise Exception("Wrong number of arguments. Following arguments expected: "
                        "<nr_of_sequences_train> <sequence_length_train> <nr_of_sequences_val> "
                        "<sequence_length_val> <nr_of_circles> <change_shape> <nr_of_epochs> <batch_size> "
                        "<ff_model_path> <lr> <lr_sched>")


import matplotlib.pyplot as plt
import matplotlib.collections as mcoll
from collections import defaultdict
from matplotlib import colors
from matplotlib.lines import Line2D

import warnings
import seaborn as sns
import math
import numpy as np
import pandas as pd

def deCasteljau(b, t):
    N = len(b)
    a = np.copy(b)
    for r in range(1, N):
        a[:N-r, :] = (1-t)*a[:N-r, :] + t*a[1:N-r+1, :]
    return a[0, :]

def BezierCv(b, nr=5):
    t = np.linspace(0, 1, nr)
    return np.array([[deCasteljau(b, t[k]),
                      deCasteljau(b, t[k+1])] for k in range(nr-1)])

def position_circle(x, radius=1):
    """ Return the x,y coordinate of the point at
        angle (360*x)°, in the circle of radius "radius"
        and center (0, 0)
    """
    return np.array([radius*math.cos(x*2*math.pi),
                     radius*math.sin(x*2*math.pi)])

def linear_gradient(start, end, n=10):
    """ Return a gradient between start and end,
        with n elements.
    """
    gradients = np.zeros((len(start), n))
    for i in range(len(start)):
        gradients[i, :] = np.linspace(start[i], end[i], num=n)
    return np.transpose(gradients)

def linear_gradient_color(c1, c2, n=10):
    """ Return a gradient between the two color c1 & c2
    """
    return linear_gradient(colors.to_rgba(c1), colors.to_rgba(c2), n=n)

def draw_chord(A, B, ax=None, color_start="b", color_end="r",
               precision=1000, **kwargs):
    """ Draw a Bezier curve between two points """
    d = np.linalg.norm(np.array(A) - np.array(B))

    # Depending on the distance between the points
    # the Bezier curve parameters change
    b = [A, A/(1 + d), B/(1 + d), B]
    bz = BezierCv(b, nr=precision)

    lc = mcoll.LineCollection(bz,
                              colors=linear_gradient_color(color_start,
                                                           color_end,
                                                           n=precision), **kwargs)
    ax.add_collection(lc)

def draw_arc_circle(start, end, color="b", radius=1, ax=None,
                    thickness=0.1, precision=1000, **kwargs):
    ts = np.linspace(start, end, precision)
    poly_nodes = ([position_circle(t, radius=radius) for t in ts] +
                  [position_circle(t, radius=radius+thickness)
                   for t in ts[::-1]])
    x, y = zip(*poly_nodes)
    ax.fill(x, y, color=color, **kwargs)

def add_text_circle(x, txt, radius=1, ax=None, **kwargs):
    """ Add text on the border of the circle, in the right orientation """
    ax.text(*position_circle(x, radius=radius),
            txt, rotation=(360*x - 180 if 0.25 < x < 0.75 else 360*x),
            ha='right' if 0.75 > x > 0.25 else 'left',
            va='top' if 0.75 > x > 0.25  else 'bottom',
            rotation_mode='anchor', **kwargs)

class Chords:
    def __init__(self, data, order_col, pair_col, color_col=None,
                 layout_args={}, text_args={},
                 chords_args={}, palette=sns.color_palette()):

        if 'spacing' not in layout_args:
            layout_args['spacing'] = 0
        if 'precision_chord' not in layout_args:
            layout_args['precision_chord'] = 100
        if 'precision_circle' not in layout_args:
            layout_args['precision_circle'] = 100
        if 'thickness_circle' not in layout_args:
            layout_args['thickness_circle'] = 0.1
        if 'subcircle' not in layout_args:
            layout_args['subcircle'] = True
        if 'radius_subcircle' not in layout_args:
            layout_args['radius_subcircle'] = 1.14
        if 'radius_circle' not in layout_args:
            layout_args['radius_circle'] = 1.02
        if 'thickness_subcircle' not in layout_args:
            layout_args['thickness_subcircle'] = 0.1
        if 'internal_chords' not in layout_args:
            layout_args['internal_chords'] = False
        if 'radius_text' not in layout_args:
            layout_args['radius_text'] = max(layout_args['thickness_subcircle']
                                             + layout_args['radius_subcircle'],
                                             layout_args['thickness_circle']
                                             + layout_args['radius_circle']) + 0.1
        if 'no_chords' not in layout_args:
            layout_args['no_chords'] = False
        if 'inverted_grad' not in layout_args:
            layout_args['inverted_grad'] = True
        if 'circle_args' not in layout_args:
            layout_args['circle_args'] = {}
        if 'subcircle_args' not in layout_args:
            layout_args['subcircle_args'] = {}
        if 'singleton' not in layout_args:
            layout_args['singleton'] = True

        if 'palette' not in text_args:
            if color_col is None:
                text_args['palette'] = palette
            else:
                text_args['palette'] = defaultdict(lambda: 'k')

        if not np.all(data[pair_col].value_counts() <= 2):
            raise TypeError("Every value in the `pair` column "
                            "should appear exactly twice")

        if not layout_args['singleton']:
            self.data = data[data.pair_col.map(data.pair_col.value_counts()) == 2]
        else:
            self.data = data.copy()

        self.chords = []
        self.data = data.copy()
        self.order_col = order_col
        self.pair_col = pair_col
        if color_col is None:
            color_col = order_col
        self.color_col = color_col
        self.df = None
        self.layout = layout_args
        self.text_args = text_args
        self.chords_args = chords_args
        self.palette = palette

        self.order_data("order", "pair")
        self.compute_positions()
        self.pair_chords()

    def order_data(self, categories, pairs):
        """
        Return a correctly ordered dataframe, ready to be plotted in chord format
        @ Args:
        data: pd.DataFrame() to reorder, with a column
        `categories` and a column `pair`
        """
        self.format_data()
        self.df["associate_cat_order"] = self.df.apply(
            lambda r: (len(self.mapcat)+r["nbcat"]-r["associate_nbcat"]) % len(self.mapcat)
            + (len(self.mapcat)//2+1 if r["nbcat"]==r["associate_nbcat"] else 0.5),
            axis=1)
        self.df["sort_order"] = self.df.apply(
            lambda r: (r["idx"] if r["nbcat"] <= r["associate_nbcat"]
                       else -r["associate"]), axis=1)
        sign = lambda x: 0 if x == 0 else 1 if x > 0 else -1
        self.df["singleton_sort"] = self.df.apply(lambda r: 0 if r["nbcat"] != r["associate_nbcat"]
                                                  else sign(r["idx"] - r["associate"]), axis=1)
        self.df["internal_sort"] = self.df.apply(lambda r: (0 if r["nbcat"] != r["associate_nbcat"]
                                      else (r["idx"] if r["idx"] < r["associate"]
                                       else -r["associate"])), axis=1)
        self.df = self.df.sort_values(by=["nbcat", "associate_cat_order", "singleton_sort",
                                          "internal_sort", "sort_order"])

    def format_data(self):
        """
        Process the dataframe so that it can be plotted in chord format
        @ Args:
        data: pd.DataFrame() to reorder, with a column
        `categories` and a column `pair`
        """
        if self.color_col == self.order_col:
            self.df = self.data[[self.order_col, self.pair_col]].rename({
                self.pair_col: "pair",
                self.order_col: "order"}, axis=1).copy()
            self.df["color"] = self.df["order"]
        else:
            self.df = self.data[[self.order_col, self.pair_col, self.color_col]].rename({
                self.pair_col: "pair",
                self.order_col: "order",
                self.color_col: "color"}, axis=1).copy()

        self.df.index.names = ['og_idx']
        self.df = self.df.reset_index()

        catunique = self.df["order"].unique()
        self.mapcat = dict(zip(catunique, range(len(catunique))))
        colorunique = self.df["color"].unique()
        self.mapcolor = dict(zip(colorunique, range(len(colorunique))))
        self.df["nbcat"] = self.df["order"].map(self.mapcat).astype(int)
        self.df["nbcolor"] = self.df["color"].map(self.mapcolor).astype(int)
        self.df["idx"] = self.df.index

        pairmap = self.df.groupby("pair").idx.apply(list).to_dict()
        self.df["associate"] = self.df.apply(
            lambda r: [a for a in pairmap[r["pair"]] if a != r["idx"]][0]
            if len(pairmap[r["pair"]]) > 1 else pairmap[r["pair"]][0],
            axis=1)

        self.df["associate_nbcat"] = self.df.associate.map(self.df.nbcat)
        self.df["associate_nbcolor"] = self.df.associate.map(self.df.nbcolor)
        self.df["catcolor"] = self.df.apply(lambda r: (r["nbcat"], r["nbcolor"]),
                                            axis=1
                                           ).astype('category').cat.codes
        self.df["associate_catcolor"] = self.df.apply(lambda r:
                                                      (r["associate_nbcat"], r["associate_nbcolor"]),
                                                      axis=1
                                                     ).astype('category').cat.codes

    def compute_positions(self):
        cat_jump = list(np.where(self.df.nbcat.values[:-1]
                                 != self.df.nbcat.values[1:])[0])
        x = 0
        positions = []
        for i in range(len(self.df)):
            positions.append(x)
            if i in cat_jump:
                x += self.layout['spacing']
            x += (1 - self.layout['spacing']*(len(cat_jump)+1))/(len(self.df))
        self.df["position"] = positions
        self.df["associate_position"] = self.df.associate.map(self.df.position)

    def pair_chords(self):
        self.chords = list(zip(
            self.df.position,
            self.df.associate_position,
            self.df.nbcolor,
            self.df.associate_nbcolor,
            self.df.nbcat,
            self.df.associate_nbcat))

        # add chord except if singleton
        self.chords = [tpl for tpl in self.chords if tpl[0] != tpl[1]]

    def add_chord(self, idx1, idx2):
        dd = self.df.set_index("og_idx")
        self.chords.append(
        (dd.loc[idx1].position, dd.loc[idx2].position,
         dd.loc[idx1].nbcolor, dd.loc[idx2].nbcolor,
         dd.loc[idx1].nbcat, dd.loc[idx2].nbcat))

    def plot(self, ax=None):
        if ax is None:
            _, ax = plt.subplots(figsize=(8, 8))
            ax.axis('off')

        nb_to_name_cat = {self.mapcat[k]:k for k in self.mapcat}
        positions = self.df.position.values
        catcolors = self.df.catcolor.values
        idxs = np.where(catcolors[:-1] != catcolors[1:])[0]
        start_categorie = [0] + list(positions[idxs+1])
        end_categorie = list(positions[idxs]) + [positions[-1]]
        cats = [self.df.nbcolor.iloc[0]] + list(self.df.nbcolor.iloc[idxs+1])

        for s, e, c in zip(start_categorie, end_categorie, cats):
            draw_arc_circle(s - 0.5/len(self.df), e + 0.5/len(self.df),
                            color=self.palette[c], ax=ax,
                            precision=self.layout['precision_circle'],
                            thickness=self.layout['thickness_circle'],
                            radius=self.layout['radius_circle'],
                            **self.layout['circle_args'])

        # the radius text should correspond to categories
        cats = self.df.nbcat.values
        idxs = np.where(cats[:-1] != cats[1:])[0]
        start_categorie = [0] + list(positions[idxs+1])
        end_categorie = list(positions[idxs]) + [positions[-1]]
        cats = [cats[0]] + list(cats[idxs+1])
        for s, e, c in zip(start_categorie, end_categorie, cats):
            add_text_circle((s + e - 1/len(self.df))/2, nb_to_name_cat[c], ax=ax,
                            color=self.text_args['palette'][c],
                            radius=self.layout['radius_text'],
                            **{k: v for k, v in self.text_args.items() if k != 'palette'})

        if self.layout['subcircle']:
            catcolors = self.df.associate_catcolor.values
            idxs = np.where(catcolors[:-1] != catcolors[1:])[0]
            start_subcategorie = [0] + list(positions[idxs+1])
            end_subcategorie = list(positions[idxs]) + [positions[-1]]
            subcats = [self.df.associate_nbcolor.iloc[0]] + list(self.df.associate_nbcolor.iloc[idxs+1])

            for s, e, c in zip(start_subcategorie, end_subcategorie, subcats):
                draw_arc_circle(s - 0.5/len(self.df), e + 0.5/len(self.df),
                                color=self.palette[c], ax=ax,
                                precision=self.layout['precision_circle'],
                                thickness=self.layout['thickness_subcircle'],
                                radius=self.layout['radius_subcircle'],
                                **self.layout['subcircle_args'])

        if not self.layout['no_chords']:
            for pos_1, pos_2, color_1, color_2, cat_1, cat_2 in self.chords:
                if cat_1 != cat_2 or self.layout['internal_chords']:
                    draw_chord(position_circle(pos_2),
                               position_circle(pos_1), ax=ax,
                               color_start=self.palette[color_2
                                                   if self.layout['inverted_grad']
                                                   else color_1],
                               color_end=self.palette[color_1
                                                 if self.layout['inverted_grad']
                                                 else color_2],
                               precision=self.layout['precision_chord'],
                               **self.chords_args)

        ax.set_xlim(-1, 1)
        ax.set_ylim(-1, 1)
        ax.axis('equal')
        ax.axis('off')
        return ax

def chord_diagram(categories, pairs, hues=None,
                  data=None, ax=None, palette=sns.color_palette(),
                  layout_args={}, text_args={}, chord_args={}):
    """ Draw a chord diagram.
        @ Args
        - categories: Categories of each individual.
        Decide the order of the plot.
        Either a list or a column name if `data` is not None.
        - pair: For each individual identifies the pair it's in.
        Every value should appear twice. Either a list or a column
        name if `data` is not None.
        - hues: list of categories that will determine the color of the plot.
        - data: dataset containing the columns `categories` and `pair`
        - ax: matplotlib ax object
        - palette: seaborn palette
        - layout_args: dict arguments for the layout, include:
            * 'spacing' (default 0): space between the categories
            * precision_chord: precision to plot the chord,
                               higher = better but slower.
            * precision_circle: same for the circles
            * subcircle: presence or not of a subcircle (see examples)
            * thickness_circle, thickness_subcircle: width of the circle / subcircle (default 0.1)
            * radius_circle, radius_subcircle: radii of both circles
            * internal_chords: Plot or not the internal chords (default `False`)
            * radius_text: radius of the text
            * no_chords: Don't plot the chords (good for testing, default `False`)
            * inverted_grad: Inverse the gradient on the chords (default `True`)
            * circle_args / subcircle_args: dict, default `{}`, additional arguement for ax.fill
            * nuplets: allow for more than one link going from the same node. Default `False`
            * singletons: allow for nodes with no "paired" value, default `True`
            * plot: Default `True`, if `False` does not plot the figure.
    """
    if 'nuplets' not in layout_args:
        layout_args['nuplets'] = False
    if 'plot' not in layout_args:
        layout_args['plot'] = True
    if layout_args['nuplets']:
        layout_args['singletons'] = True

    if data is None:
        data = pd.DataFrame()
        data["cat"] = categories
        data["pair"] = pairs
        data["col"] = hues
        categories = "cat"
        pairs = "pair"
        hues = None if hues is None else "col"

    doublets = None
    if layout_args['nuplets']:
        data_copy = data.copy()
        data_copy.index.names = ['idx']
        data_copy = data_copy.reset_index()
        map_pair = data_copy.groupby(pairs).idx.apply(list)
        new_pairs = [str(p) + str(map_pair[p].index(idx)//2)
                    for idx, p in zip(data_copy.idx, data_copy[pairs])]
        data_copy[pairs] = new_pairs
    else:
        data_copy = data.copy()

    ch = Chords(data=data_copy, order_col=categories,
                pair_col=pairs,
                color_col=hues,
                 layout_args=layout_args, text_args=text_args,
                 chords_args=chord_args, palette=palette)

    if layout_args['nuplets']:
        for idx1, p1, p2 in zip(data.index, data[pairs], data_copy[pairs]):
            if len(map_pair) > 2:
                for idx2 in map_pair[p1]:
                    ch.add_chord(idx1, idx2)

    if layout_args['plot']:
        ch.plot(ax=ax)
    return ch


import os
import csv

#describe the path to find the files needed
csv_file_path = os.path.join("election_data.csv")

#identify variables
total_votes = 0         # count total number of votes
candidates = {}         # dictionary of keys (candidate names) and values (number of votes)
candidates_percent = {} # dictionary of vote percentages for candidates
winner_count = 0        # number of votes for winning candidate
winner_name = ""        #name of the winning candidate in an empty string

#------------------------------------------------------------------------------

#with open(csv_file_path, newLine="") as csv_file:
with open(csv_file_path, newline="") as csv_file:
   reader = csv.reader(csv_file, delimiter = ",")
   next(reader, None) #check the next method
   for row in reader:
       total_votes += 1 # tally total vote count

       # find list of candidates who received votes
       if row[2] in candidates.keys():
           candidates[row[2]] += 1
       else:
           candidates[row[2]] = 1
       #we are counting the candidates and building the dictionary of candidates where the key is the candidate
       #name and the value is the number of votes.
       # calculate percentages for votes won by each candidate (candidates dictionary: key= Name, value= number of votes)
       for key, value in candidates.items():
           candidates_percent[key] = round((value/total_votes) * 100, 1)
            #items is a sequence in the dictionary that returns you key value pairs
       # find winning candidate using candidates dictionary
       #round *100, 1 will make a percentage of the fraction and keep one decimal point

       for key in candidates.keys():
           if candidates[key] > winner_count:
               winner_name = key
               winner_count = candidates[key]

#Format to print out the results
print ("Election Results")
print ("------------------------")
print ("Total Votes: " + str(total_votes))
print("-------------------------")
for key, value in candidates.items():
    print(key + ":" + str(candidates_percent[key]) + "% (" + str(value) + ")")
print("--------------------------")
print("Winner: " + winner_name)
print("-------------------------")

output_file = "results.txt"
with open(output_file, 'w') as outfile:
    outfile.write("Election Results \n")
    outfile.write("--------------------------\n")
    outfile.write("Total Votes: " + str(total_votes) + "\n")
    outfile.write("-------------------------\n")

   #for key, value in candidates.items():
for key, value in candidates.items():
       outfile.write(key + ": " + str(candidates_percent[key]) + "% (" + str(value) + ") \n")
       outfile.write("--------------------------\n")
       outfile.write("Winner: " + winner_name + "\n")
       outfile.write("-------------------------\n")


import turtle

wn = turtle.Screen()
wn.title("Estrella tortuga Python")
wn.bgcolor("lightgreen")

leonardo = turtle.Turtle()
leonardo.hideturtle()
leonardo.pensize(3)
leonardo.fillcolor("#ffa41b")
leonardo.color("#ffa41b")

side = wn.numinput("Tamaño", "Indica un tamaño en píxeles para la estrella")

leonardo.begin_fill()
for i in range (5):
    leonardo.forward(side)
    leonardo.right(144)
leonardo.end_fill()

leonardo.write("Estrella", False, "right", ("arial", int(side/4), "bold italic"))

wn.exitonclick()


'''
Description: Definitions of functions
Author: Nan Li
Contact: linan.lqq0@gmail.com
'''
# import packages
import re
import time
import pandas as pd
from headers import headers
from datetime import datetime
from seleniumwire import webdriver

# short names dictionary
address_dict = ({'road':'rd', 'street':'st', 'place':'pl', 'avenue':'ave',
                'parade':'pde', 'highway':'hwy', 'drive':'dr', 'grove':'gr',
                'crescent':'cres', 'court':'ct', 'close':'cl', 'circuit':'cct',
                'avenue':'ave',
                })

# function to read postcode from excel file
def readPostcode(file_name):
    postcode = pd.read_excel(file_name).sort_values(by=['Postcode'])
    postcode = postcode.values
    postcode = postcode[:,0].astype(int)
    return postcode

# function to create valid url according to postcode
def createUrl(postcode, page_no):
    url_pre = 'https://www.realestate.com.au/sold/in-'
    if postcode < 1000:
        url = '{}0{}/list-{}?includeSurrounding=false'.format(url_pre, postcode, page_no)
    elif postcode < 10000:
        url = '{}{}/list-{}?includeSurrounding=false'.format(url_pre, postcode, page_no)
    else:
        print('Wrong postcode!!')
        exit()
    return url

# function to create valid url according to address
def createHouseUrl1(address, suburb, postcode, use_short):
    url_pre = 'https://www.realestate.com.au/property/'
    str_address = address.lower()
    if re.search(r'[0-9|a-z]+/[0-9]', str_address) is not None:
        str_address = 'unit-' + str_address
    if use_short:
        for key in address_dict.keys():
            street_name = re.search(key, str_address)
            if street_name is not None:
                str_address = re.sub(street_name.group(), address_dict[key], str_address)
    str_suburb = suburb.lower()
    if postcode < 1000:
        str_postcode = 'nt-0{:d}'.format(postcode)
    elif postcode < 2600:
        str_postcode = 'nsw-{:d}'.format(postcode)
    elif postcode < 2700:
        str_postcode = 'act-{:d}'.format(postcode)
    elif postcode < 2800:
        str_postcode = 'nsw-{:d}'.format(postcode)
    elif postcode < 3000:
        str_postcode = 'act-{:d}'.format(postcode)
    elif postcode < 4000:
        str_postcode = 'vic-{:d}'.format(postcode)
    elif postcode < 5000:
        str_postcode = 'qld-{:d}'.format(postcode)
    elif postcode < 6000:
        str_postcode = 'sa-{:d}'.format(postcode)
    elif postcode < 7000:
        str_postcode = 'wa-{:d}'.format(postcode)
    elif postcode < 8000:
        str_postcode = 'tas-{:d}'.format(postcode)
    str_address = '{}-{}-{}'.format(str_address, str_suburb, str_postcode)
    str_address = re.sub('[,|\s|/]', '-', str_address)
    for _ in range(4):
        str_address = re.sub('--', '-', str_address)
    url = url_pre + str_address
    return url

# function to create url according to REA_id
def createHouseUrl2(REA_id):
    url = 'https://www.realestate.com.au/property/lookup?id={:d}'.format(REA_id)
    return url

def get_cookie(url, user_agent):
    options = webdriver.ChromeOptions()
    options.add_argument('user-agent={}'.format(user_agent))
    options.add_argument("--window-size=100x100")
    options.add_argument('ignore-certificate-errors')
    options.add_argument("--disable-blink-features=AutomationControlled")
    options.add_experimental_option("excludeSwitches", ["enable-automation", "enable-logging"])
    try:
        driver = webdriver.Chrome(options=options)
    except:
        print('Update ChromeDriver')
    '''driver._orig_get = driver.get
    def _get_wrapped(*args, **kwargs):
        if driver.execute_script("return navigator.webdriver"):
            driver.execute_cdp_cmd(
                "Page.addScriptToEvaluateOnNewDocument",
                {
                    "source": """
                                Object.defineProperty(window, 'navigator', {
                                    value: new Proxy(navigator, {
                                    has: (target, key) => (key === 'webdriver' ? false : key in target),
                                    get: (target, key) =>
                                        key === 'webdriver'
                                        ? undefined
                                        : typeof target[key] === 'function'
                                        ? target[key].bind(target)
                                        : target[key]
                                    })
                                });
                            """
                },
            )
        return driver._orig_get(*args, **kwargs)
    driver.get = _get_wrapped
    driver.get = _get_wrapped
    driver.get = _get_wrapped
    original_user_agent_string = driver.execute_script(
        "return navigator.userAgent"
    )
    driver.execute_cdp_cmd(
        "Network.setUserAgentOverride",
        {
            "userAgent": original_user_agent_string.replace("Headless", ""),
        },
    )
    driver.execute_cdp_cmd(
        "Page.addScriptToEvaluateOnNewDocument",
        {
            "source": """
                            Object.defineProperty(navigator, 'maxTouchPoints', {
                                get: () => 1
                        })"""
        },
    )'''

    def interceptor(request):
        del request.headers['user-agent']
        request.headers['user-agent'] = user_agent
        # Block PNG, JPEG and GIF images
        if request.path.endswith(('.png', '.jpg', '.gif')):
            request.abort()

    #driver.request_interceptor = interceptor
    driver.get(url)
    driver.execute_script("window.open('{}');".format(url))
    time.sleep(4) # wait for cookies loaded
    cnt = 0
    for request in driver.requests:
        if request.response is not None and request.response.status_code == 200:
            cookie = re.search(r"(KP_UIDz-ssn=[\S]+);", str(request.response.headers))
            if cookie is not None:
                cookie = cookie.group(1)
                cnt = cnt + 1
                if cnt > 3:
                    break
    if cookie is None:
        print('Change cookie name and retry!')

    driver.quit()

    return cookie

def update_cookie(old, new):
    with open("headers.py", "r") as f:
        data = f.read()
        data = re.sub(old, new, data)

    with open("headers.py", "w") as f:
        f.write(data)


def parentheses(s):
    l=[]
    for i in range(len(s)):
        if s[i]=='(' or s[i]=='{' or s[i]=='[':
            l.append(s[i])
        else:
            if len(l)==0:
                return False
            x=l[-1]
            if x=='(' and s[i]==')':
                l.pop()
            elif x=='{' and s[i]=='}':
                l.pop()
            elif x=='[' and s[i]==']':
                l.pop()
            else:
                return False
    if len(l)!=0:
        return False
    return True
s=input()
print(parentheses(s))


import requests
import bs4
from bs4 import BeautifulSoup
import pandas as pd
import time
import csv
import simplejson as json
import threading

def extract_job_title_from_result(soup):
    jobs = []
    for div in soup.find_all(name="div", attrs={"class":"row"}):
        for a in div.find_all(name="a", attrs={"data-tn-element":"jobTitle"}):
            jobs.append(a["title"])
    return(jobs)

def extract_company_from_result(soup):
  companies = []
  for div in soup.find_all(name="div", attrs={"class":"row"}):
    company = div.find_all(name="span", attrs={"class":"company"})
    if len(company) > 0:
      for b in company:
        companies.append(b.text.strip())
    else:
      sec_try = div.find_all(name="span", attrs={"class":"result-link-source"})
      for span in sec_try:
          companies.append(span.text.strip())
  return(companies)

def extract_location_from_result(soup):
  locations = []
  spans = soup.findAll("span", attrs={"class": "location"})
  for span in spans:
    locations.append(span.text)
  return(locations)

def extract_salary_from_result(soup):
  salaries = []
  for div in soup.find_all(name="div", attrs={"class":"row"}):
    try:
      salaries.append(div.find('nobr').text)
    except:
      try:
        div_two = div.find(name="div", attrs={"class":"sjcl"})
        div_three = div_two.find("div")
        salaries.append(div_three.text.strip())
      except:
        salaries.append("Nothing_found")
  return(salaries)

def extract_summary_from_result(soup):
  summaries = []
  spans = soup.findAll("span", attrs={"class": "summary"})
  for span in spans:
    summaries.append(span.text.strip())
  return(summaries)

URL = "https://www.indeed.com/jobs?q=data+scientist+%2420%2C000&l=New+York&start=10"

page = requests.get(URL)
soup = BeautifulSoup(page.text, "html.parser")

print(soup.prettify())

if __name__ == "__main__":

    t1 = threading.Thread(target=extract_job_title_from_result, args=(soup,))
    t2 = threading.Thread(target=extract_company_from_result, args=(soup,))
    t3 = threading.Thread(target=extract_location_from_result,args=(soup,))
    t4 = threading.Thread(target=extract_salary_from_result,args=(soup,))
    t5 = threading.Thread(target=extract_summary_from_result,args=(soup,))

    t1.start()
    t2.start()
    t3.start()
    t4.start()
    t5.start()

    t1.join()
    t2.join()
    t3.join()
    t4.join()
    t5.join()

    page = requests.get(URL,timeout=5)  # + '&start=' + str(start))
    time.sleep(1)
    soup = BeautifulSoup(page.text, 'lxml')
    job_post = []
    for div in soup.find_all(name='div', attrs={'class':'row'}):
        job_post_object = {
                            "job_title": div.find(name="a").text.encode('utf-8'),
                            "company": div.find(name="span").text.encode('utf-8'),
                            "location": div.find(name="span").text.encode('utf-8'),
                            "summary": div.find(name='span').text.encode('utf-8'),
                            "salary": div.find(name="div").text.encode('utf-8')
                    }
        job_post.append(job_post_object)
    with open('IndeedData.json', 'w') as outfile:
        json.dump(job_post, outfile)
        outfile.write('\n')


import json

#build the dict, takes in the file name as param
def buildDict(s):
  fp = open(s, "r")
  dictionary = {}
  linesInFile = []
  for line in fp:
    newLine = line.split('@')
    linesInFile.append(newLine)

  length = len(linesInFile)
  for i in range(0,length):
    #new innerDict for using the insidedict
    innerDict = {}
    x = 1
    y = 2
    lenOfLine = len(linesInFile[i])-1
    while( y < lenOfLine):
      #create a dict out of a new word
      innerDict.update({linesInFile[i][x]: linesInFile[i][y]})
      x = x + 2
      y = y + 2
    #add the new dict to the overall dict
    dictionary.update({linesInFile[i][0]: innerDict})

  return dictionary

# uncomment this to print the dict nicely
'''
  print ( json.dumps(dictionary, indent=2))

'''


#!/usr/bin/env python2
# engine='python'.
# -*- coding: utf-8 -*-
"""
Created on Tue Oct 25 16:43:06 2016

@author: andrepignata
"""
import urllib.request
from os import getcwd

import pandas as pd
import numpy as np

cwd = getcwd()
print(cwd)

aCidades = pd.read_csv('bases/cidades.csv')
#aTemas = [{'ano':2010,'tema':1}]
aTemas = [{'ano':2015,'tema':118}]
rCidade = aCidades.ix[0]
iTema = aTemas[0]
for rCidade in aCidades:
    for iTema in aTemas:
        sArquivo = './csv/'+str(rCidade[0])+'_'+str(iTema['tema'])+'.csv'
        print(sArquivo)
        urllib.request.urlretrieve('http://cidades.ibge.gov.br/xtras/csv.php?lang=&idtema='+str(iTema['tema'])+'&codmun='+str(rCidade[0]),sArquivo)
        with open(sArquivo, 'r',encoding='iso-8859-1') as f:
            content = f.readlines()
        iEspaco = 0;
        aDado = list()
        for sLinha in content:
            if (sLinha.strip() == ''):
                iEspaco += 1
            if iEspaco == 1:
                sLinha = sLinha.replace(';-;',';Null;')
                sLinha = sLinha.replace('Não informado','Null')
                sLinha = sLinha.replace('.','')
                sLinha = sLinha.replace(',','.')
                aDado.append(sLinha.split(';'))
        aDado.pop(0)
        aDado.pop(0)

        data = pd.DataFrame(aDado)

        data['ano'] = iTema['ano']
        data['codmun'] = rCidade[0]
        data.columns = ['descricao','valor','unidade','ano','codmun']
        data.replace('Null',np.nan,inplace=True)
        print(data)
        #data.to_sql('idh',engine,if_exists='append')


def buddyStrings(s, goal):
    if len(s) != len(goal):
        return False

    if s == goal:
        # Check if there are at least two repeated characters in s
        char_count = {}
        for char in s:
            if char in char_count:
                return True
            char_count[char] = 1
        return False

    indices = []
    for i in range(len(s)):
        if s[i] != goal[i]:
            indices.append(i)

    if len(indices) != 2:
        return False

    i, j = indices
    return s[i] == goal[j] and s[j] == goal[i]

# Test the example input
s = "ab"
goal = "ba"
result = buddyStrings(s, goal)
print(result)  # Output: True


import matplotlib.pyplot as plt
import glob
import os
import argparse

# File select index out of range exception
class ValueOutOfRange(Exception):
    def __str__(self):
        return 'Value out of range'

data_path = None

parser = argparse.ArgumentParser()
parser.add_argument('--data', '-d', nargs='?', action='store', dest='data_path', metavar='data file')
parser.add_argument('--version', '-v', action='version', version='v1.0')
args = parser.parse_args()

try:
    if args.data_path is None:
        data_path =  os.path.dirname(os.path.abspath(__file__)) + '/../data/*.TXT'
    else:
        sel = 0
        data_path = args.data_path

    # Find for data in data/ file
    data_files = glob.glob(data_path)

    # List found files
    if args.data_path is None:
        for i, fn in enumerate(data_files):
            print('[', i + 1, '] ', os.path.basename(fn))

        # Select file
        sel = None
        while sel not in range(len(data_files)) or type(sel) != int:
            try:
                sel = int(input('Select file > ')) - 1

                # Exit if entered value is 0
                if sel == -1:
                    exit()

                if sel not in range(len(data_files)):
                    raise ValueOutOfRange()
            except ValueOutOfRange as e:
                print(e)
            except ValueError as e:
                print(e)

    # Open data file
    try:
        f = open(data_files[sel], 'r')
    except Exception as e:
        print(e)

    # HYPERPARAMETERS
    DATA_LINE = 3
    ZERO_LINE = 0

    # Parse data
    try:
        for i, line in enumerate(f):
            if i == DATA_LINE:
                data = line.split(':')
            if i == ZERO_LINE:
                zero_line = line.split(':')

        f.close()  # Close data file
    except Exception as e:
        print(e)

    f.close()  # Close data file

    # Parse tare value
    for item in zero_line:
        try:
            tare = int(item)
        except:
            pass

    print("Tare, zero value: ", tare)

    # Convert to integers and append to int_data list
    int_data = data
    for i, item in enumerate(data):
        try:
            int_data[i] = int(data[i])
        except:
            del int_data[i]

    del [int_data[len(int_data) - 1]]

    # Apply tare
    for i, d in enumerate(int_data):
        int_data[i] = abs(d) - abs(tare)  # With tare
        int_data[i] /= 20000  # *20000 to convert to Kh¡g and *9.8 for Newtopns

    plt.plot(range(len(int_data)), int_data, 'y')
    plt.grid()
    plt.xlabel("Lecture")
    plt.ylabel("Kgf")
    plt.show()
except KeyboardInterrupt as e:
    print("\nBye!")