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()


from db import model
from api import api
from datetime import datetime
import json
import webbrowser
from vkinder.settings import Settings

class APP:

    def __init__(self, token, _id):
        self.vk = api.VK(token)
        self.user_id = str(self.vk.resolve_screen_name(_id))
        self.db = model.DB(self.user_id)
        self.settings = Settings(self.user_id)

    @staticmethod
    def user_input():
        url = 'https://oauth.vk.com/authorize?client_id=7331062&display=page&redirect_uri=https://oauth.vk.com/blank.html&scope=friends,photos,groups&response_type=token&v=5.103'
        webbrowser.open_new_tab(url)
        token = input('Введите токен: ')
        _id = input('Введите ваш id: ')
        return APP(token, _id)

    @staticmethod
    def progress_bar(data, func, name, one=False, bar_len=50):
        start = datetime.now()
        data_len = len(data)
        result = []
        for i, value in enumerate(data):
            p_bar = round(i / data_len * bar_len)
            bar = f'[{"*"*p_bar + " "*(bar_len - p_bar)}]'
            print(
                f'{name:^15} {bar} {i} of {data_len} {datetime.now() - start}\r',
                end=''
                )
            if one:
                func(value)
            else:
                result.extend(func(**value))
        print(f'{name} [{"*"*bar_len}] {data_len} of {data_len} {datetime.now() -start} done')
        return result

    def get_top_three_profile_photos(self, _id):
        photos = self.vk.get_top_three_profile_photos(_id)
        if not photos:
            return 'Профиль закрыт или без фото'

        def parse(photo):
            return {
                'w': 6,
                'z': 5,
                'y': 4,
                'x': 3,
                'm': 2,
                's': 1,
            }.get(photo['type'], 0)

        return [max(photo['sizes'], key=parse)['url'] for photo in photos]

    def out(self, users):
        file_name = f'out_{self.user_id}.json'
        with open(file_name, mode='w', encoding='utf-8') as file:
            json.dump(
                [
                    {f'https://vk.com/{user["domain"]}':
                        self.get_top_three_profile_photos(user['id'])}
                    for user in users
                ],
                file, ensure_ascii=False, indent=4)
            print(f'результат записан в {file_name}')

    def check_settings(self):
        obligatory_search_params = [
            'age_from', 'age_to', 'sex', 'country', 'city'
        ]
        for param in obligatory_search_params:
            if param not in self.settings.search:
                self.settings.search[param] =\
                    input(f'Нехватает {param}. Введите значение: ')
        if 'group_id' not in self.settings.search:
            load = input('Введите id групп через запятую или "load" для загрузки.')
            if load == 'load':
                self.settings.search['group_id'] =\
                    self.vk.get_groups(self.user_id)
            else:
                self.settings.search['group_id'] = load.split(', ')


import requests
import os

CHAT_ID = os.getenv('CHAT_ID')
TOKEN = os.getenv('TG_TOKEN')

print(f"TOKEN is {TOKEN}, CHAT_ID is {CHAT_ID}")

if not CHAT_ID or not TOKEN:
    raise ValueError("Char ID or TOKEN was not specified")

TOKEN = TOKEN.replace("\n", "")

async def send_message(data: str):
    """
    Send data to the telegrame bot
    :param data: Data to send
    :return: True if sent successfully
    """
    send_text = 'https://api.telegram.org/bot' + TOKEN + '/sendMessage?chat_id=' \
                + CHAT_ID + '&parse_mode=Markdown&text=' + data
    print(f"Send text was {send_text}")
    response = requests.get(send_text)
    return response.ok


import pickle
from flask import Flask, request, app, jsonify,url_for, render_template
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler

app=Flask(__name__)
## load the model
regmodel=pickle.load(open('regmodel.pkl','rb'))
scalar=pickle.load(open('scaling.pkl','rb'))
@app.route('/')
def home():
    return render_template('home.html')
@app.route('/predict_api',methods=['POST'])
def predict_api():
    data=request.json['data']
    print(data)
    ## the data values receieved must be transformed into a list which is then changed into np array
    ## the numpy array is then reshaped to a single row, multiple columns(1,13)
    print(np.array(list(data.values())).reshape(1,-1))
    new_data=scalar.transform(np.array(list(data.values())).reshape(1,-1))
    output=regmodel.predict(new_data)
    # since this is a two dimensional array, we need output[0]
    print(output[0])
    return jsonify(output[0])

if __name__=="__main__":
    app.run(debug=True)


import os
import json
import sqlite3
import datetime, time
import itertools
from common import util
import queue
import threading
from threading import Thread
import logging

import sqlite3
import datetime, time
import itertools
from common import util

class IDEBenchDriver:

    def init(self, options, schema, driver_arg):
        self.time_of_latest_request = 0
        self.isRunning = False
        self.requests = queue.LifoQueue()
        self.config = json.load(open(os.path.join(os.path.dirname(__file__),'..','sqlite.config.json')))

    def create_connection(self):
        sqlite_file = self.config['dbFilename']
        conn = sqlite3.connect(sqlite_file)
        return conn

    def sqlitefix(self, sql_statement):
        if not "FLOOR" in sql_statement:
            return sql_statement
        else:
            sql_statement=sql_statement.replace("FLOOR", "ROUND")
            x=sql_statement.find("ROUND")
            y=sql_statement.find(")",x)
            output=sql_statement[:y]+" -0.5 "+sql_statement[y:]
            #print(output,flush=True)
            return output

    def execute_vizrequest(self, viz_request, options, schema, result_queue):

        viz = viz_request.viz
        sql_statement = viz.get_computed_filter_as_sql(schema)
        #calculate connection time

        connection = self.conn
        cursor = connection.cursor()

        viz_request.start_time = util.get_current_ms_time()
        #print(sql_statement,flush=True,end = '')

        cursor.execute(self.sqlitefix(sql_statement))
        data = cursor.fetchall()
        viz_request.end_time = util.get_current_ms_time()

        cursor.close()

        results = {}
        for row in data:
            keys = []
            for i, bin_desc in enumerate(viz.binning):

                if "width" in bin_desc:
                    bin_width = bin_desc["width"]
                    keys.append(str(int(row[i])))
                else:
                    keys.append(str(row[i]))

            key = ",".join(keys)
            results[key] = row[len(viz.binning):]

        viz_request.result = results
        result_queue.put(viz_request)

    def process_request(self, viz_request, options, schema, result_queue):
        self.requests.put((viz_request, options, schema, result_queue))

    def process(self):
        self.conn = self.create_connection()

        while self.isRunning:
            try:
                request = self.requests.get(timeout=1)
                viz_request = request[0]
                options = request[1]
                schema = request[2]
                result_queue = request[3]

                # only execute requests that are newer than the last one we processed (drops old/no longer needed queries)
                if viz_request.expected_start_time < self.time_of_latest_request:
                    viz_request.dropped = True
                    result_queue.put(viz_request)
                    continue

                self.time_of_latest_request = viz_request.expected_start_time
                self.execute_vizrequest(viz_request, options, schema, result_queue)
            except Exception as e:
                # ignore queue-empty exceptions
                pass

        # close connection when done
        self.conn.close()

    def workflow_start(self):
        self.isRunning = True
        thread = Thread(target = self.process)
        thread.start()

    def workflow_end(self):
        self.isRunning = False


import json

ores = [
    'iron',
    'coal',
    'gold',
    'copper',
    'silver',
    'redstone',
    'diamond',
    'lapis',
    'emerald',
    'quartz',
    'tin',
    'lead',
    'nickel',
    'zinc',
    'aluminum',
    'cobalt',
    'osmium',
    'iridium',
    'uranium',
    'ruby',
    'sapphire',
    'sulfur',
    'cinnabar',
    'potassium_nitrate',
    'apatite',
    'fluorite'
]

stones = [
    'andesite',
    'basalt',
    'blackstone',
    'calcite',
    'deepslate',
    'diorite',
    'end_stone',
    'granite',
    'netherrack',
    'stone',
    'tuff',
    'create_limestone',
    'create_scoria',
    'create_scorchia',
    'quark_jasper',
    'quark_limestone',
    'quark_shale'
]

forge_ores = {
    "replace": False,
    "values" : []
}

forge_ores_specific = {
    "raplace": False,
    "values": []
}

for ore in ores :
    for stone in stones :
        forge_ores["values"].append('omniores:' + ore + '_ore_' + stone)
        forge_ores_specific["values"].append('omniores:' + ore + '_ore_' + stone)
    block_forge_ores_specific = open('F:/Creative/diversityMods/omniores/src/main/resources/data/forge/tags/blocks/ores/'+ore+'.json', 'w')
    block_forge_ores_specific.write(json.dumps(forge_ores_specific, indent=2))
    block_forge_ores_specific.close()
    item_forge_ores_specific = open('F:/Creative/diversityMods/omniores/src/main/resources/data/forge/tags/items/ores/'+ore+'.json', 'w')
    item_forge_ores_specific.write(json.dumps(forge_ores_specific, indent=2))
    item_forge_ores_specific.close()
    forge_ores_specific['values'] = []

block_forge_ores_file = open('F:/Creative/diversityMods/omniores/src/main/resources/data/forge/tags/blocks/ores.json', 'w')
block_forge_ores_file.write(json.dumps(forge_ores, indent=2))
block_forge_ores_file.close()

item_forge_ores_file = open('F:/Creative/diversityMods/omniores/src/main/resources/data/forge/tags/items/ores.json', 'w')
item_forge_ores_file.write(json.dumps(forge_ores, indent=2))
item_forge_ores_file.close()


from django.shortcuts import redirect, render
from . import models

def index(req):
  return render(req, 'home.html')

def listKategori(req):
  if req.method == 'POST':
    kategori_id = req.POST.get('kategori_id')
    data = models.KategoriModel.objects.get(id=kategori_id)
    data.delete()
    return redirect('/kategori/')

  return render(req, 'kategori/list-kategori.html', context={
    'kategories': models.KategoriModel.objects.all()
  })

def addKategori(req):
  if req.method == 'POST':
    nama  = req.POST.get('nama')

    data = models.KategoriModel()
    data.nama = nama
    data.save()
    return redirect('/kategori/')

  return render(req, 'kategori/add-form.html')

def editKategori(req, id):
  if req.method == 'POST':
    nama = req.POST.get('nama')
    data = models.KategoriModel.objects.get(id=id)
    data.nama = nama
    data.save()
    return redirect('/kategori/')

  return render(req, 'kategori/edit-form.html', context={
    'kategori': models.KategoriModel.objects.get(id=id)
  })

def listToko(req):
  if req.method == 'POST':
    toko_id = req.POST.get('toko_id')
    data = models.TokoModel.objects.get(id=toko_id)
    data.delete()
    return redirect('/toko/')

  return render(req, 'toko/list-toko.html', context={
    'tokos': models.TokoModel.objects.all()
  })

def addToko(req):
  if req.method == 'POST':
    nama    = req.POST.get('nama')
    pemilik = req.POST.get('pemilik')
    email   = req.POST.get('email')
    no_telp = req.POST.get('no_telp')
    logo    = req.FILES.get('logo')
    alamat  = req.POST.get('alamat')

    data = models.TokoModel()
    data.nama     = nama
    data.pemilik  = pemilik
    data.email    = email
    data.no_telp  = no_telp
    data.logo     = logo
    data.alamat   = alamat
    data.save()
    return redirect('/toko/')

  return render(req, 'toko/add-form.html')

def editToko(req, id):
  if req.method == 'POST':
    data = models.TokoModel.objects.get(id=id)

    if req.FILES:
      logo      = req.FILES.get('logo')
      data.logo = logo

    nama    = req.POST.get('nama')
    pemilik = req.POST.get('pemilik')
    email   = req.POST.get('email')
    no_telp = req.POST.get('no_telp')
    alamat  = req.POST.get('alamat')

    data.nama     = nama
    data.pemilik  = pemilik
    data.email    = email
    data.no_telp  = no_telp
    data.alamat   = alamat
    data.save()
    return redirect('/toko/')

  return render(req, 'toko/edit-form.html', context={
    'toko': models.TokoModel.objects.get(id=id)
  })

def listProduct(req):
  if req.method == 'POST':
    product_id = req.POST.get('product_id')
    data = models.ProdukModel.objects.get(id=product_id)
    data.delete()
    return redirect('/product/')

  return render(req, 'product/list-product.html', context={
    'products': models.ProdukModel.objects.all()
  })

def addProduct(req):
  if req.method == 'POST':
    nama      = req.POST.get('nama')
    harga     = req.POST.get('harga')
    berat     = req.POST.get('berat')
    stok      = req.POST.get('stok')
    gambar    = req.FILES.get('gambar')
    kategori  = req.POST.getlist('kategori')
    deskripsi = req.POST.get('deskripsi')

    data = models.ProdukModel()
    data.nama = nama
    data.harga = harga
    data.berat = berat
    data.stok = stok
    data.gambar = gambar
    data.deskripsi= deskripsi
    data.save()
    for i in kategori:
      data.kategori.add(i)
    return redirect('/products/')

  return render(req, 'product/add-form.html', context={
    'kategories': models.KategoriModel.objects.all()
  })

def editProduct(req, id):
  if req.method == 'POST':

    data = models.ProdukModel.objects.get(id=id)

    if req.FILES:
      gambar    = req.FILES.get('gambar')
      data.gambar = gambar

    nama      = req.POST.get('nama')
    harga     = req.POST.get('harga')
    berat     = req.POST.get('berat')
    stok      = req.POST.get('stok')
    kategori  = req.POST.getlist('kategori')
    deskripsi = req.POST.get('deskripsi')

    data.nama = nama
    data.harga = harga
    data.berat = berat
    data.stok = stok
    data.deskripsi= deskripsi
    data.save()
    for i in kategori:
      data.kategori.add(i)
    return redirect('/product/')

  productKategories = models.ProdukModel.objects.filter(id=id).values('kategori')
  productKategoriID = []
  for i in productKategories:
    productKategoriID.append(i.get('kategori'))

  return render(req, 'product/edit-form.html',context={
    'product': models.ProdukModel.objects.get(id=id),
    'kategories': models.KategoriModel.objects.all(),
    'productKategori': productKategoriID
  })


import cv2
from mtcnn import MTCNN

capture = cv2.VideoCapture(0)

detector = MTCNN()

while True:

    ret, frame = capture.read()
    faces = detector.detect_faces(frame)

    for single_faces in faces:
        x, y, width, height = single_faces["box"]

        left_eyeX, left_eyeY = single_faces["keypoints"]["left_eye"]
        right_eyeX, right_eyeY = single_faces["keypoints"]["right_eye"]
        noseX, noseY = single_faces["keypoints"]["nose"]
        mouth_leftX, mouth_leftY = single_faces["keypoints"]["mouth_left"]
        mouth_rightX, mouth_rightY = single_faces["keypoints"]["mouth_right"]

        cv2.rectangle(frame, pt1=(x, y), pt2=(x + width, y + height), color=(255, 0, 0), thickness=3)
        cv2.circle(frame, center=(left_eyeX, left_eyeY), color=(255, 0, 0), thickness=3, radius=2)
        cv2.circle(frame, center=(right_eyeX, right_eyeY), color=(255, 0, 0), thickness=3, radius=2)
        cv2.circle(frame, center=(noseX, noseY), color=(255, 0, 0), thickness=3, radius=2)
        cv2.circle(frame, center=(mouth_leftX, mouth_leftY), color=(255, 0, 0), thickness=3, radius=2)
        cv2.circle(frame, center=(mouth_rightX, mouth_rightY), color=(255, 0, 0), thickness=3, radius=2)
        cv2.rectangle(frame, pt1=(x, y), pt2=(x + width, y + height), color=(255, 0, 0), thickness=3)

    cv2.imshow("front_cam", frame)

    if cv2.waitKey(1) & 0xFF == ord("b"):
        break

cv2.destroyAllWindows()


from flask import render_template,request,redirect
from models import *
from forms import CardForm, StoryForm, NewsForm
from upload import save_file
from app import app
from extensions import db
import requests
import xml.etree.ElementTree as ET

@app.route('/')
def index():
    story = Story.query.all()
    card = Cards.query.all()
    currency = ForeignCurrency.query.all()
    news = News.query.all()
    return render_template('index.html',story=story,cards=card,currency=currency,news=news)

@app.route('/cards/')
def card():
    card = Cards.query.all()
    story = Story.query.all()
    return render_template('card.html',cards=card,story=story, page_name='cards')

@app.route("/loans/")
def credits():
    card = Cards.query.all()
    story = Story.query.all()
    return render_template('kredit.html',page_name='credits',cards=card,story=story,)

@app.route('/add-card/', methods=['GET', 'POST'])
def addcard():
    form = CardForm()

    if request.method == 'POST':
        print('post')
        if form.validate:
            card_title = form.title.data
            card_desc = form.description.data
            card_term = form.term.data
            card_currency = form.currency.data
            card_cashback = form.cashback.data
            image_data = form.image.data
            print(image_data)
            card_image = save_file(image_data)
            card = Cards(card_title,card_desc,card_term,card_currency,card_cashback,card_image)
            card.save()
            print(card)
            return redirect('/cards')
    return render_template('card_add.html', form=form)

@app.route('/add-story/', methods=['GET', 'POST'])
def addStory():
    form = StoryForm()

    if request.method == 'POST':
        if form.validate_on_submit():
            story_tip = form.tip.data
            story_description = form.description.data
            story_image = save_file(form.image.data)
            story_color = form.color.data
            storys = Story(story_tip,story_description,story_image,story_color)
            storys.save()
            return redirect('/cards')
    return render_template('story_add.html', form=form)

@app.route('/valyuta/')
def loadRSS():
    url = f'https://www.cbar.az/currencies/{datetime.datetime.now().strftime("%d.%m.%Y")}.xml'
    response = requests.get(url, stream=True)
    tree = ET.fromstring(response.content)
    dom = tree.findall('ValType')

    db.session.query(ForeignCurrency).delete()
    db.session.commit()

    for attr_type in dom:
        name = attr_type.get('Type')
        if name == 'Xarici valyutalar':
            valute = attr_type.findall('Valute')
            for code in valute:
                code_name = code.get('Code')
                if code_name == 'USD' or code_name == "AZN" or code_name == "EUR":
                    nominal = code.find('Nominal').text
                    name = code.find('Name').text
                    course = code.find('Value').text
                    currency = ForeignCurrency(code_name, nominal, name, course)
                    currency.save()
    return redirect('/cards')

@app.route('/add-news/', methods=['GET', 'POST'])
def addnews():
    form = NewsForm()

    if request.method == 'POST':
        if form.validate_on_submit():
            news_title = form.title.data
            news = News(news_title)
            news.save()
    return render_template('news_add.html', form=form)


from mirl.agents.td3_agent import TD3Agent
import mirl.torch_modules.utils as ptu
import torch.nn.functional as F
import copy
from mirl.utils.logger import logger

class TD3BCAgent(TD3Agent):
    def __init__(
        self,
        env,
        policy,
        qf,
        qf_target,
        pool=None,
        normalize_obs=True,
        bc_alpha=2.5,
        **td3_kwargs
    ):
        super().__init__(
            env,
            policy,
            qf,
            qf_target,
            **td3_kwargs
        )
        self.bc_alpha = bc_alpha
        self.normalize_obs = normalize_obs
        if normalize_obs:
            assert pool is not None
            assert pool.compute_mean_std
            self.obs_mean, self.obs_std = pool.get_mean_std()['observations']
            self.obs_mean = ptu.from_numpy(self.obs_mean)
            self.obs_std = ptu.from_numpy(self.obs_std)

    def process_obs(self, obs):
        if self.normalize_obs:
            obs = obs = (obs-self.obs_mean)/(self.obs_std+1e-6)
        return obs

    def step_explore(self, o, **kwargs):
        raise NotImplementedError

    def step_exploit(self, o, **kwargs):
        with self.policy.deterministic_(True):
            o = ptu.from_numpy(o)
            o = self.process_obs(o)
            a,_ = self.policy.action(o, **kwargs)
            a = ptu.get_numpy(a)
        return a, {}

    def compute_policy_loss(
        self, obs, new_action,
        origin_action, v_pi_kwargs={}
    ):
        q_new_action, _ = self.qf.value(
            obs,
            new_action,
            return_ensemble=False,
            **v_pi_kwargs
        )
        q_pi_mean = q_new_action.mean()
        policy_loss = -q_pi_mean
        policy_info = self._log_policy_info(
            new_action, policy_loss, q_pi_mean)
        ### for td3bc ###
        bc_lambda = self.bc_alpha / q_pi_mean.detach().abs().mean()
        bc_loss = F.mse_loss(new_action, origin_action)
        policy_loss = policy_loss*bc_lambda + bc_loss
        if self._log_tb_or_not():
            logger.tb_add_scalar("policy/bc_loss", bc_loss, self.num_train_steps)
        ###### end ######
        return policy_loss, policy_info

    def train_from_torch_batch(self, batch):
        rewards = batch['rewards']
        terminals = batch['terminals']
        obs = batch['observations']
        actions = batch['actions']
        next_obs = batch['next_observations']
        obs = self.process_obs(obs)
        next_obs = self.process_obs(next_obs)
        self.log_batch(rewards, terminals)
        ################
        # update critic #
        ################
        q_target = self.compute_q_target(next_obs, rewards, terminals, self.next_v_pi_kwargs)
        qf_loss, train_qf_info = self.compute_qf_loss(obs, actions, q_target)
        self.qf_optimizer.zero_grad()
        qf_loss.backward()
        self.log_critic_grad_norm()
        self.qf_optimizer.step()
        if self.num_train_steps % self.target_update_freq == 0:
            self._update_target(self.soft_target_tau)
        self.train_info.update(train_qf_info)
        ################
        # update actor #
        ################
        if self.num_train_steps % self.policy_update_freq == 0:
            new_action, action_info = self.policy.action(obs, **self.current_sample_kwargs)
            policy_loss, train_policy_info = self.compute_policy_loss(
                obs,
                new_action,
                actions,  #NOTE: origin action for bc
                self.current_v_pi_kwargs
            )
            self.policy_optimizer.zero_grad()
            policy_loss.backward()
            self.policy_optimizer.step()
            self.train_info.update(train_policy_info)
        #####################
        # update statistics #
        #####################
        if self._need_to_update_eval_statistics:
            self._need_to_update_eval_statistics = False
            self.eval_statistics.update(self.train_info)
        self.log_train_info()
        self.num_train_steps += 1
        return copy.deepcopy(self.train_info)


# coding=utf8
from youtubesearchpython import Search
import pandas as pd
import numpy as np

search_query_test = 'Кончится Лето Kino'

def find_first_vid(search_query):

    allSearch = Search(search_query, limit = 1)

    result_dict = allSearch.result()
    result_dict2 = list(result_dict.values())[-1]
    result_dict3 = result_dict2[-1]
    result_dict4 = list(result_dict3.values())[-2]

    return result_dict4

df_songs = pd.read_csv('song_links_results.csv', encoding='utf8')
df_list_songs = df_songs.values.tolist()
song_list = []

for i in df_list_songs:
    song_list.append(i[1])

from pytube import YouTube
import os
from pathlib import Path

def youtube2mp3 (url,outdir):
    # url input from user
    yt = YouTube(url)

    ##@ Extract audio with 160kbps quality from video
    video = yt.streams.filter(abr='160kbps').last()

    ##@ Downloadthe file
    out_file = video.download(output_path=outdir)
    base, ext = os.path.splitext(out_file)
    new_file = Path(f'{base}.mp3')
    os.rename(out_file, new_file)
    ##@ Check success of download
    if new_file.exists():
        print(f'{yt.title} has been successfully downloaded.')
    else:
        print(f'ERROR: {yt.title}could not be downloaded!')

'''
for song in song_list[1619:]:
    print(song_list.index(song))
    try:
        youtube2mp3(song,'dirpath')
    except FileExistsError:
        continue
    except AttributeError:
        continue
    except KeyError:
        continue

df_convert = pd.read_csv('converted_spotify_music_list.csv', encoding='utf8')
#df.apply(lambda x: pd.lib.infer_dtype(x.values))

df_titles_artists = df_convert[['Track Name','Artist Name(s)']]
df_titles_artists.astype(str)
df_titles_artists["separator"] = '-'

df_titles_artists["combined_col"] = df_titles_artists["Track Name"] + df_titles_artists["separator"] + df_titles_artists["Artist Name(s)"]

df_combined = df_titles_artists["combined_col"]

print(df_combined)
df_combined = df_combined.dropna(axis=0)
titles_artists_list = df_combined.values.tolist()

#print(titles_artists_list)

song_link_list = []
#titles_artists_list_subset = titles_artists_list[0:50]
titles_artists_list_subset = titles_artists_list
for song in titles_artists_list_subset:
    print('trying:',song)
    result_link = find_first_vid(song)
    song_link_list.append(result_link)

df_song_links = pd.DataFrame(song_link_list)

df_song_links.to_csv('song_links_results.csv', encoding='utf-8')

'''