from jax.config import config

config.update("jax_enable_x64", True)
import jax

import pytest
import numpy as np

from rdkit import Chem
from rdkit.Chem import AllChem
from ff.handlers import nonbonded, bonded
from ff.handlers.deserialize import deserialize_handlers

import functools

def test_harmonic_bond():

    patterns = [
        ["[#6X4:1]-[#6X4:2]", 0.1, 0.2],
        ["[#6X4:1]-[#6X3:2]", 99.0, 99.0],
        ["[#6X4:1]-[#6X3:2]=[#8X1+0]", 99.0, 99.0],
        ["[#6X3:1]-[#6X3:2]", 99.0, 99.0],
        ["[#6X3:1]:[#6X3:2]", 99.0, 99.0],
        ["[#6X3:1]=[#6X3:2]", 99.0, 99.0],
        ["[#6:1]-[#7:2]", 0.1, 0.2],
        ["[#6X3:1]-[#7X3:2]", 99.0, 99.0],
        ["[#6X4:1]-[#7X3:2]-[#6X3]=[#8X1+0]", 99.0, 99.0],
        ["[#6X3:1](=[#8X1+0])-[#7X3:2]", 99.0, 99.0],
        ["[#6X3:1]-[#7X2:2]", 99.0, 99.0],
        ["[#6X3:1]:[#7X2,#7X3+1:2]", 99.0, 99.0],
        ["[#6X3:1]=[#7X2,#7X3+1:2]", 99.0, 99.0],
        ["[#6:1]-[#8:2]", 99.0, 99.0],
        ["[#6X3:1]-[#8X1-1:2]", 99.0, 99.0],
        ["[#6X4:1]-[#8X2H0:2]", 0.3, 0.4],
        ["[#6X3:1]-[#8X2:2]", 99.0, 99.0],
        ["[#6X3:1]-[#8X2H1:2]", 99.0, 99.0],
        ["[#6X3a:1]-[#8X2H0:2]", 99.0, 99.0],
        ["[#6X3:1](=[#8X1])-[#8X2H0:2]", 99.0, 99.0],
        ["[#6:1]=[#8X1+0,#8X2+1:2]", 99.0, 99.0],
        ["[#6X3:1](~[#8X1])~[#8X1:2]", 99.0, 99.0],
        ["[#6X3:1]~[#8X2+1:2]~[#6X3]", 99.0, 99.0],
        ["[#6X2:1]-[#6:2]", 99.0, 99.0],
        ["[#6X2:1]-[#6X4:2]", 99.0, 99.0],
        ["[#6X2:1]=[#6X3:2]", 99.0, 99.0],
        ["[#6:1]#[#7:2]", 99.0, 99.0],
        ["[#6X2:1]#[#6X2:2]", 99.0, 99.0],
        ["[#6X2:1]-[#8X2:2]", 99.0, 99.0],
        ["[#6X2:1]-[#7:2]", 99.0, 99.0],
        ["[#6X2:1]=[#7:2]", 99.0, 99.0],
        ["[#16:1]=[#6:2]", 99.0, 99.0],
        ["[#6X2:1]=[#16:2]", 99.0, 99.0],
        ["[#7:1]-[#7:2]", 99.0, 99.0],
        ["[#7X3:1]-[#7X2:2]", 99.0, 99.0],
        ["[#7X2:1]-[#7X2:2]", 99.0, 99.0],
        ["[#7:1]:[#7:2]", 99.0, 99.0],
        ["[#7:1]=[#7:2]", 99.0, 99.0],
        ["[#7+1:1]=[#7-1:2]", 99.0, 99.0],
        ["[#7:1]#[#7:2]", 99.0, 99.0],
        ["[#7:1]-[#8X2:2]", 99.0, 99.0],
        ["[#7:1]~[#8X1:2]", 99.0, 99.0],
        ["[#8X2:1]-[#8X2:2]", 99.0, 99.0],
        ["[#16:1]-[#6:2]", 99.0, 99.0],
        ["[#16:1]-[#1:2]", 99.0, 99.0],
        ["[#16:1]-[#16:2]", 99.0, 99.0],
        ["[#16:1]-[#9:2]", 99.0, 99.0],
        ["[#16:1]-[#17:2]", 99.0, 99.0],
        ["[#16:1]-[#35:2]", 99.0, 99.0],
        ["[#16:1]-[#53:2]", 99.0, 99.0],
        ["[#16X2,#16X1-1,#16X3+1:1]-[#6X4:2]", 99.0, 99.0],
        ["[#16X2,#16X1-1,#16X3+1:1]-[#6X3:2]", 99.0, 99.0],
        ["[#16X2:1]-[#7:2]", 99.0, 99.0],
        ["[#16X2:1]-[#8X2:2]", 99.0, 99.0],
        ["[#16X2:1]=[#8X1,#7X2:2]", 99.0, 99.0],
        ["[#16X4,#16X3!+1:1]-[#6:2]", 99.0, 99.0],
        ["[#16X4,#16X3:1]~[#7:2]", 99.0, 99.0],
        ["[#16X4,#16X3:1]-[#8X2:2]", 99.0, 99.0],
        ["[#16X4,#16X3:1]~[#8X1:2]", 99.0, 99.0],
        ["[#15:1]-[#1:2]", 99.0, 99.0],
        ["[#15:1]~[#6:2]", 99.0, 99.0],
        ["[#15:1]-[#7:2]", 99.0, 99.0],
        ["[#15:1]=[#7:2]", 99.0, 99.0],
        ["[#15:1]~[#8X2:2]", 99.0, 99.0],
        ["[#15:1]~[#8X1:2]", 99.0, 99.0],
        ["[#16:1]-[#15:2]", 99.0, 99.0],
        ["[#15:1]=[#16X1:2]", 99.0, 99.0],
        ["[#6:1]-[#9:2]", 99.0, 99.0],
        ["[#6X4:1]-[#9:2]", 0.6, 0.7],
        ["[#6:1]-[#17:2]", 99.0, 99.0],
        ["[#6X4:1]-[#17:2]", 99.0, 99.0],
        ["[#6:1]-[#35:2]", 99.0, 99.0],
        ["[#6X4:1]-[#35:2]", 99.0, 99.0],
        ["[#6:1]-[#53:2]", 99.0, 99.0],
        ["[#6X4:1]-[#53:2]", 99.0, 99.0],
        ["[#7:1]-[#9:2]", 99.0, 99.0],
        ["[#7:1]-[#17:2]", 99.0, 99.0],
        ["[#7:1]-[#35:2]", 99.0, 99.0],
        ["[#7:1]-[#53:2]", 99.0, 99.0],
        ["[#15:1]-[#9:2]", 99.0, 99.0],
        ["[#15:1]-[#17:2]", 99.0, 99.0],
        ["[#15:1]-[#35:2]", 99.0, 99.0],
        ["[#15:1]-[#53:2]", 99.0, 99.0],
        ["[#6X4:1]-[#1:2]", 99.0, 99.0],
        ["[#6X3:1]-[#1:2]", 99.0, 99.0],
        ["[#6X2:1]-[#1:2]", 99.0, 99.0],
        ["[#7:1]-[#1:2]", 99.0, 99.0],
        ["[#8:1]-[#1:2]", 99.0, 99.1],
    ]

    smirks = [x[0] for x in patterns]
    params = np.array([[x[1], x[2]] for x in patterns])
    props = None
    hbh = bonded.HarmonicBondHandler(smirks, params, props)

    mol = Chem.MolFromSmiles("C1CNCOC1F")

    bond_params, bond_idxs = hbh.parameterize(mol)

    assert bond_idxs.shape == (mol.GetNumBonds(), 2)
    assert bond_params.shape == (mol.GetNumBonds(), 2)

    bonded_param_adjoints = np.random.randn(*bond_params.shape)

    bond_params_new, bond_vjp_fn, bond_idxs_new = jax.vjp(
        functools.partial(hbh.partial_parameterize, mol=mol), hbh.params, has_aux=True
    )

    np.testing.assert_array_equal(bond_params_new, bond_params)
    np.testing.assert_array_equal(bond_idxs_new, bond_idxs)

    # test that we can use the adjoints
    ff_adjoints = bond_vjp_fn(bonded_param_adjoints)[0]

    # if a parameter is > 99 then its adjoint should be zero (converse isn't necessarily true since)
    mask = np.argwhere(bond_params > 90)
    assert np.all(ff_adjoints[mask] == 0.0) == True

def test_proper_torsion():

    # proper torsions have a variadic number of terms

    patterns = [
        ["[*:1]-[#6X3:2]=[#6X3:3]-[*:4]", [[99.0, 99.0, 99.0]]],
        ["[*:1]-[#6X3:2]=[#6X3:3]-[#35:4]", [[99.0, 99.0, 99.0]]],
        ["[#9:1]-[#6X3:2]=[#6X3:3]-[#35:4]", [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]],
        ["[#35:1]-[#6X3:2]=[#6X3:3]-[#35:4]", [[7.0, 8.0, 9.0], [1.0, 3.0, 5.0], [4.0, 4.0, 4.0]]],
        ["[#9:1]-[#6X3:2]=[#6X3:3]-[#9:4]", [[7.0, 8.0, 9.0]]],
    ]

    smirks = [x[0] for x in patterns]
    params = [x[1] for x in patterns]
    props = None

    pth = bonded.ProperTorsionHandler(smirks, params, props)

    mol = Chem.MolFromSmiles("FC(Br)=C(Br)F")

    torsion_params, torsion_idxs = pth.parameterize(mol)

    assert torsion_idxs.shape == (8, 4)
    assert torsion_params.shape == (8, 3)

    torsion_params_new, torsion_vjp_fn, torsion_idxs_new = jax.vjp(
        functools.partial(pth.partial_parameterize, mol=mol), pth.params, has_aux=True
    )

    np.testing.assert_array_equal(torsion_params_new, torsion_params)
    np.testing.assert_array_equal(torsion_idxs_new, torsion_idxs)

    torsion_param_adjoints = np.random.randn(*torsion_params.shape)

    ff_adjoints = torsion_vjp_fn(torsion_param_adjoints)[0]

    mask = np.argwhere(torsion_params > 90)
    assert np.all(ff_adjoints[mask] == 0.0) == True

def test_improper_torsion():

    patterns = [
        ["[*:1]~[#6X3:2](~[*:3])~[*:4]", 1.5341333333333333, 3.141592653589793, 2.0],
        ["[*:1]~[#6X3:2](~[#8X1:3])~[#8:4]", 99.0, 99.0, 99.0],
        ["[*:1]~[#7X3$(*~[#15,#16](!-[*])):2](~[*:3])~[*:4]", 99.0, 99.0, 99.0],
        ["[*:1]~[#7X3$(*~[#6X3]):2](~[*:3])~[*:4]", 1.3946666666666667, 3.141592653589793, 2.0],
        ["[*:1]~[#7X3$(*~[#7X2]):2](~[*:3])~[*:4]", 99.0, 99.0, 99.0],
        ["[*:1]~[#7X3$(*@1-[*]=,:[*][*]=,:[*]@1):2](~[*:3])~[*:4]", 99.0, 99.0, 99.0],
        ["[*:1]~[#6X3:2](=[#7X2,#7X3+1:3])~[#7:4]", 99.0, 99.0, 99.0],
    ]

    smirks = [x[0] for x in patterns]
    params = np.array([[x[1], x[2], x[3]] for x in patterns])
    props = None
    imp_handler = bonded.ImproperTorsionHandler(smirks, params, props)

    mol = Chem.MolFromSmiles("CNC(C)=O")  # peptide
    mol = Chem.AddHs(mol)

    torsion_params, torsion_idxs = imp_handler.parameterize(mol)

    assert torsion_idxs.shape[0] == 6  # we expect two sets of impropers, each with 3 components.
    assert torsion_idxs.shape[1] == 4

    assert torsion_params.shape[0] == 6
    assert torsion_params.shape[1] == 3

    torsion_params_new, torsion_vjp_fn, torsion_idxs_new = jax.vjp(
        functools.partial(imp_handler.partial_parameterize, mol=mol), imp_handler.params, has_aux=True
    )

    np.testing.assert_array_equal(torsion_params_new, torsion_params)
    np.testing.assert_array_equal(torsion_idxs_new, torsion_idxs)

    param_adjoints = np.random.randn(*torsion_params.shape)

    # # test that we can use the adjoints
    ff_adjoints = torsion_vjp_fn(param_adjoints)[0]

    # # if a parameter is > 99 then its adjoint should be zero (converse isn't necessarily true since)
    mask = np.argwhere(torsion_params > 90)
    assert np.all(ff_adjoints[mask] == 0.0) == True

def test_exclusions():

    mol = Chem.MolFromSmiles("FC(F)=C(F)F")
    exc_idxs, scales = nonbonded.generate_exclusion_idxs(mol, scale12=0.0, scale13=0.2, scale14=0.5)

    for pair, scale in zip(exc_idxs, scales):
        src, dst = pair
        assert src < dst

    expected_idxs = np.array(
        [
            [0, 1],
            [0, 2],
            [0, 3],
            [0, 4],
            [0, 5],
            [1, 2],
            [1, 3],
            [1, 4],
            [1, 5],
            [2, 3],
            [2, 4],
            [2, 5],
            [3, 4],
            [3, 5],
            [4, 5],
        ]
    )

    np.testing.assert_equal(exc_idxs, expected_idxs)

    expected_scales = [0.0, 0.2, 0.2, 0.5, 0.5, 0.0, 0.0, 0.2, 0.2, 0.2, 0.5, 0.5, 0.0, 0.0, 0.2]
    np.testing.assert_equal(scales, expected_scales)

def test_am1_bcc():
    # currently takes no parameters
    smirks = []
    params = []
    props = None

    am1h = nonbonded.AM1BCCHandler(smirks, params, props)
    mol = Chem.AddHs(Chem.MolFromSmiles("C1CNCOC1F"))
    AllChem.EmbedMolecule(mol)
    charges = am1h.parameterize(mol)

    assert len(charges) == mol.GetNumAtoms()

    new_charges, vjp_fn = jax.vjp(functools.partial(am1h.partial_parameterize, mol=mol))

    # charges_adjoints = np.random.randn(*charges.shape)

    # assert vjp_fn(charges_adjoints) == None

def test_am1_ccc():

    patterns = [
        ["[#6X4:1]-[#1:2]", 0.46323257920556493],
        ["[#6X3$(*=[#8,#16]):1]-[#6a:2]", 0.24281402370571598],
        ["[#6X3$(*=[#8,#16]):1]-[#8X1,#8X2:2]", 1.0620166764992722],
        ["[#6X3$(*=[#8,#16]):1]=[#8X1$(*=[#6X3]-[#8X2]):2]", 2.227759732057297],
        ["[#6X3$(*=[#8,#16]):1]=[#8X1,#8X2:2]", 2.8182928673804217],
        ["[#6a:1]-[#8X1,#8X2:2]", 0.5315976926761063],
        ["[#6a:1]-[#1:2]", 0.0],
        ["[#6a:1]:[#6a:2]", 0.0],
        ["[#6a:1]:[#6a:2]", 0.0],
        ["[#8X1,#8X2:1]-[#1:2]", -2.3692047944101415],
        ["[#16:1]-[#8:2]", 99.0],
    ]

    smirks = [x[0] for x in patterns]
    params = np.array([x[1] for x in patterns])
    props = None
    mol_sdf = """
     RDKit          3D

 21 21  0  0  0  0  0  0  0  0999 V2000
    1.9775   -1.1921   -0.0420 O   0  0  0  0  0  0  0  0  0  0  0  0
    2.0482   -0.0099    0.4175 C   0  0  0  0  0  0  0  0  0  0  0  0
    3.2666    0.4539    1.1431 C   0  0  0  0  0  0  0  0  0  0  0  0
    0.9469    0.7961    0.2044 O   0  0  0  0  0  0  0  0  0  0  0  0
   -0.1834    0.3655   -0.4688 C   0  0  0  0  0  0  0  0  0  0  0  0
   -0.3446    0.4891   -1.8261 C   0  0  0  0  0  0  0  0  0  0  0  0
   -1.4607    0.0654   -2.5014 C   0  0  0  0  0  0  0  0  0  0  0  0
   -2.4866   -0.5199   -1.7874 C   0  0  0  0  0  0  0  0  0  0  0  0
   -2.3616   -0.6611   -0.4259 C   0  0  0  0  0  0  0  0  0  0  0  0
   -1.2246   -0.2274    0.2480 C   0  0  0  0  0  0  0  0  0  0  0  0
   -1.1499   -0.4031    1.6843 C   0  0  0  0  0  0  0  0  0  0  0  0
   -2.1061   -0.9387    2.2814 O   0  0  0  0  0  0  0  0  0  0  0  0
   -0.0450    0.0107    2.3913 O   0  0  0  0  0  0  0  0  0  0  0  0
    3.6820    1.3873    0.6660 H   0  0  0  0  0  0  0  0  0  0  0  0
    3.0604    0.6688    2.2138 H   0  0  0  0  0  0  0  0  0  0  0  0
    4.0231   -0.3405    1.1004 H   0  0  0  0  0  0  0  0  0  0  0  0
    0.4612    0.9512   -2.4027 H   0  0  0  0  0  0  0  0  0  0  0  0
   -1.5273    0.1914   -3.5726 H   0  0  0  0  0  0  0  0  0  0  0  0
   -3.3631   -0.8543   -2.3065 H   0  0  0  0  0  0  0  0  0  0  0  0
   -3.1913   -1.1293    0.1168 H   0  0  0  0  0  0  0  0  0  0  0  0
   -0.0216    0.8969    2.8664 H   0  0  0  0  0  0  0  0  0  0  0  0
  1  2  2  0
  2  3  1  0
  2  4  1  0
  4  5  1  0
  5  6  2  0
  6  7  1  0
  7  8  2  0
  8  9  1  0
  9 10  2  0
 10 11  1  0
 11 12  2  0
 11 13  1  0
 10  5  1  0
  3 14  1  0
  3 15  1  0
  3 16  1  0
  6 17  1  0
  7 18  1  0
  8 19  1  0
  9 20  1  0
 13 21  1  0
M  END
$$$$
    """
    am1h = nonbonded.AM1CCCHandler(smirks, params, props)
    mol = Chem.MolFromMolBlock(mol_sdf, removeHs=False)
    es_params = am1h.parameterize(mol)

    # TBD update with AM1 Symmetrize=True
    ligand_params = np.array(
        [
            -6.10948,
            7.08286,
            -1.13097,
            -4.3096,
            2.03822,
            -1.72492,
            -0.98493,
            -1.75663,
            -0.70852,
            -1.65444,
            7.76463,
            -6.49881,
            -7.10585,
            0.93707,
            0.93707,
            0.93707,
            1.79376,
            1.6654,
            1.69322,
            1.90986,
            5.22498,
        ]
    )

    np.testing.assert_almost_equal(es_params, ligand_params, decimal=5)

    new_es_params, es_vjp_fn = jax.vjp(functools.partial(am1h.partial_parameterize, mol=mol), params)

    np.testing.assert_array_equal(es_params, new_es_params)

    es_params_adjoints = np.random.randn(*es_params.shape)

    adjoints = es_vjp_fn(es_params_adjoints)[0]

    # if a parameter is > 99 then its adjoint should be zero (converse isn't necessarily true since)
    mask = np.argwhere(params > 90)
    assert np.all(adjoints[mask] == 0.0) == True

    import time

    start = time.time()
    es_params_from_cache = am1h.parameterize(mol)
    end = time.time()

    # second pass should be very fast
    assert end - start < 1.0

    # should be *exactly* identical since we're loading from cache
    np.testing.assert_array_equal(es_params_from_cache, es_params)

def test_simple_charge_handler():

    patterns = [
        ["[#1:1]", 99.0],
        ["[#1:1]-[#6X4]", 99.0],
        ["[#1:1]-[#6X4]-[#7,#8,#9,#16,#17,#35]", 99.0],
        ["[#1:1]-[#6X4](-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]", 99.0],
        ["[#1:1]-[#6X4](-[#7,#8,#9,#16,#17,#35])(-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]", 99.0],
        ["[#1:1]-[#6X4]~[*+1,*+2]", 99.0],
        ["[#1:1]-[#6X3]", 99.0],
        ["[#1:1]-[#6X3]~[#7,#8,#9,#16,#17,#35]", 99.0],
        ["[#1:1]-[#6X3](~[#7,#8,#9,#16,#17,#35])~[#7,#8,#9,#16,#17,#35]", 99.0],
        ["[#1:1]-[#6X2]", 99.0],
        ["[#1:1]-[#7]", 99.0],
        ["[#1:1]-[#8]", 99.0],
        ["[#1:1]-[#16]", 99.0],
        ["[#6:1]", 0.7],
        ["[#6X2:1]", 99.0],
        ["[#6X4:1]", 0.1],
        ["[#8:1]", 99.0],
        ["[#8X2H0+0:1]", 0.5],
        ["[#8X2H1+0:1]", 99.0],
        ["[#7:1]", 0.3],
        ["[#16:1]", 99.0],
        ["[#15:1]", 99.0],
        ["[#9:1]", 1.0],
        ["[#17:1]", 99.0],
        ["[#35:1]", 99.0],
        ["[#53:1]", 99.0],
        ["[#3+1:1]", 99.0],
        ["[#11+1:1]", 99.0],
        ["[#19+1:1]", 99.0],
        ["[#37+1:1]", 99.0],
        ["[#55+1:1]", 99.0],
        ["[#9X0-1:1]", 99.0],
        ["[#17X0-1:1]", 99.0],
        ["[#35X0-1:1]", 99.0],
        ["[#53X0-1:1]", 99.0],
    ]

    smirks = [x[0] for x in patterns]
    params = np.array([x[1] for x in patterns])
    props = None

    sch = nonbonded.SimpleChargeHandler(smirks, params, props)

    mol = Chem.MolFromSmiles("C1CNCOC1F")

    es_params = sch.parameterize(mol)

    ligand_params = np.array([0.1, 0.1, 0.3, 0.1, 0.5, 0.1, 1.0])  # C  # C  # N  # C  # O  # C  # F

    np.testing.assert_almost_equal(es_params, ligand_params)

    es_params_adjoints = np.random.randn(*es_params.shape)

    new_es_params, es_vjp_fn = jax.vjp(functools.partial(sch.partial_parameterize, mol=mol), params)

    np.testing.assert_array_equal(new_es_params, es_params)
    # test that we can use the adjoints
    adjoints = es_vjp_fn(es_params_adjoints)[0]

    # if a parameter is > 99 then its adjoint should be zero (converse isn't necessarily true since)
    mask = np.argwhere(params > 90)
    assert np.all(adjoints[mask] == 0.0) == True

@pytest.mark.skip("gbsa is deprecated")
def test_gbsa_handler():

    patterns = [
        ["[*:1]", 99.0, 99.0],
        ["[#1:1]", 99.0, 99.0],
        ["[#1:1]~[#7]", 99.0, 99.0],
        ["[#6:1]", 0.1, 0.2],
        ["[#7:1]", 0.3, 0.4],
        ["[#8:1]", 0.5, 0.6],
        ["[#9:1]", 0.7, 0.8],
        ["[#14:1]", 99.0, 99.0],
        ["[#15:1]", 99.0, 99.0],
        ["[#16:1]", 99.0, 99.0],
        ["[#17:1]", 99.0, 99.0],
    ]

    smirks = [x[0] for x in patterns]
    params = np.array([[x[1], x[2]] for x in patterns])
    props = {"foo": "bar"}
    gbh = nonbonded.GBSAHandler(smirks, params, props)

    mol = Chem.MolFromSmiles("C1CNCOC1F")

    gb_params, gb_vjp_fn = gbh.parameterize(mol)

    ligand_params = np.array(
        [
            [0.1, 0.2],  # C
            [0.1, 0.2],  # C
            [0.3, 0.4],  # N
            [0.1, 0.2],  # C
            [0.5, 0.6],  # O
            [0.1, 0.2],  # C
            [0.7, 0.8],  # F
        ]
    )

    np.testing.assert_almost_equal(gb_params, ligand_params)

    gb_params_adjoints = np.random.randn(*gb_params.shape)

    # test that we can use the adjoints
    adjoints = gb_vjp_fn(gb_params_adjoints)[0]

    # if a parameter is > 99 then its adjoint should be zero (converse isn't necessarily true since)
    mask = np.argwhere(params > 90)
    assert np.all(adjoints[mask] == 0.0) == True

def test_am1_differences():

    ff_raw = open("ff/params/smirnoff_1_1_0_ccc.py").read()
    ff_handlers = deserialize_handlers(ff_raw)
    for ccc in ff_handlers:
        if isinstance(ccc, nonbonded.AM1CCCHandler):
            break

    suppl = Chem.SDMolSupplier("tests/data/ligands_40.sdf", removeHs=False)
    smi = "[H]c1c(OP(=S)(OC([H])([H])C([H])([H])[H])OC([H])([H])C([H])([H])[H])nc(C([H])(C([H])([H])[H])C([H])([H])[H])nc1C([H])([H])[H]"
    smi = "Clc1c(Cl)c(Cl)c(-c2c(Cl)c(Cl)c(Cl)c(Cl)c2Cl)c(Cl)c1Cl"
    mol = Chem.MolFromSmiles(smi)
    mol = Chem.AddHs(mol)
    mol.SetProp("_Name", "Debug")
    assert AllChem.EmbedMolecule(mol) == 0

    suppl = [mol]
    am1 = nonbonded.AM1Handler([], [], None)
    bcc = nonbonded.AM1BCCHandler([], [], None)

    for mol in suppl:

        am1_params = am1.parameterize(mol)
        ccc_params = ccc.parameterize(mol)
        bcc_params = bcc.parameterize(mol)

        if np.sum(np.abs(ccc_params - bcc_params)) > 0.1:

            print(mol.GetProp("_Name"), Chem.MolToSmiles(mol))
            print("  AM1    CCC    BCC  S ?")
            for atom_idx, atom in enumerate(mol.GetAtoms()):
                a = am1_params[atom_idx]
                b = bcc_params[atom_idx]
                c = ccc_params[atom_idx]
                print("{:6.2f}".format(a), "{:6.2f}".format(c), "{:6.2f}".format(b), atom.GetSymbol(), end="")
                if np.abs(b - c) > 0.1:
                    print(" *")
                else:
                    print(" ")

            assert 0

def test_lennard_jones_handler():

    patterns = [
        ["[#1:1]", 99.0, 999.0],
        ["[#1:1]-[#6X4]", 99.0, 999.0],
        ["[#1:1]-[#6X4]-[#7,#8,#9,#16,#17,#35]", 99.0, 999.0],
        ["[#1:1]-[#6X4](-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]", 99.0, 999.0],
        ["[#1:1]-[#6X4](-[#7,#8,#9,#16,#17,#35])(-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]", 99.0, 999.0],
        ["[#1:1]-[#6X4]~[*+1,*+2]", 99.0, 999.0],
        ["[#1:1]-[#6X3]", 99.0, 999.0],
        ["[#1:1]-[#6X3]~[#7,#8,#9,#16,#17,#35]", 99.0, 999.0],
        ["[#1:1]-[#6X3](~[#7,#8,#9,#16,#17,#35])~[#7,#8,#9,#16,#17,#35]", 99.0, 999.0],
        ["[#1:1]-[#6X2]", 99.0, 999.0],
        ["[#1:1]-[#7]", 99.0, 999.0],
        ["[#1:1]-[#8]", 99.0, 999.0],
        ["[#1:1]-[#16]", 99.0, 999.0],
        ["[#6:1]", 0.7, 0.8],
        ["[#6X2:1]", 99.0, 999.0],
        ["[#6X4:1]", 0.1, 0.2],
        ["[#8:1]", 99.0, 999.0],
        ["[#8X2H0+0:1]", 0.5, 0.6],
        ["[#8X2H1+0:1]", 99.0, 999.0],
        ["[#7:1]", 0.3, 0.4],
        ["[#16:1]", 99.0, 999.0],
        ["[#15:1]", 99.0, 999.0],
        ["[#9:1]", 1.0, 1.1],
        ["[#17:1]", 99.0, 999.0],
        ["[#35:1]", 99.0, 999.0],
        ["[#53:1]", 99.0, 999.0],
        ["[#3+1:1]", 99.0, 999.0],
        ["[#11+1:1]", 99.0, 999.0],
        ["[#19+1:1]", 99.0, 999.0],
        ["[#37+1:1]", 99.0, 999.0],
        ["[#55+1:1]", 99.0, 999.0],
        ["[#9X0-1:1]", 99.0, 999.0],
        ["[#17X0-1:1]", 99.0, 999.0],
        ["[#35X0-1:1]", 99.0, 999.0],
        ["[#53X0-1:1]", 99.0, 999.0],
    ]

    smirks = [x[0] for x in patterns]
    params = np.array([[x[1], x[2]] for x in patterns])
    props = None

    ljh = nonbonded.LennardJonesHandler(smirks, params, props)

    mol = Chem.MolFromSmiles("C1CNCOC1F")

    lj_params = ljh.parameterize(mol)

    ligand_params = np.array(
        [
            [0.1 / 2, 0.2],  # C
            [0.1 / 2, 0.2],  # C
            [0.3 / 2, 0.4],  # N
            [0.1 / 2, 0.2],  # C
            [0.5 / 2, 0.6],  # O
            [0.1 / 2, 0.2],  # C
            [1.0 / 2, 1.1],  # F
        ]
    )

    np.testing.assert_almost_equal(lj_params, ligand_params)

    lj_params_adjoints = np.random.randn(*lj_params.shape)

    new_lj_params, lj_vjp_fn = jax.vjp(functools.partial(ljh.partial_parameterize, mol=mol), params)

    np.testing.assert_array_equal(new_lj_params, lj_params)

    # test that we can use the adjoints
    adjoints = lj_vjp_fn(lj_params_adjoints)[0]

    # if a parameter is > 99 then its adjoint should be zero (converse isn't necessarily true since)
    mask = np.argwhere(params > 90)
    assert np.all(adjoints[mask] == 0.0) == True


import tkinter
from tkinter import *
import cv2
import PIL.Image, PIL.ImageTk
import time
import argparse
import os
from keras import backend as K
import tensorflow as tf
from scipy.spatial import distance as dist
from imutils.video import VideoStream
from imutils import face_utils
from threading import Thread
import numpy as np
import playsound
import argparse
import imutils
import time
import dlib
import cv2
import pickle
import pyttsx3
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from collections import deque

from src.vehicles.number_plate_recognition.number_plate_recognition_service import NumberPlateRecognitionService

class SpeedLimitsViolationService:

	def __init__(self, args):

		# modules

		self.number_plate_recognition_service = NumberPlateRecognitionService()

		# variables

		self.violators_queue = deque(maxlen=3)
		self.last_time_violated = dict()
		self.num_violations = 0

	def find_by_id(self, detection_queue, car_id):

		for obj in detection_queue:
			if obj['id'] == car_id:
				return obs

		return None

	def process(self, prev_ego_state, ego_state, frame_queue, detection_queue, current_car_ids, curr_time, dt, max_speed):

		new_violation_candidates = set()
		result_violators = set()

		# find violators

		for car_id in current_car_ids:

			if len(detection_queue) < 2 or self.find_by_id(detection_queue[-2], car_id) is None:
				continue

			prev_car_state = self.find_by_id(detection_queue[-2], car_id)
			curr_car_state = self.find_by_id(detection_queue[-1], car_id)

			prev_car_x = prev_car_state['center'][0] + prev_ego_state['position_x']
			prev_car_y = prev_car_state['center'][1] + prev_ego_state['position_y']

			curr_car_x = curr_car_state['center'][0] + ego_state['position_x']
			curr_car_y = curr_car_state['center'][1] + ego_state['position_y']

			pos_shift_x, pos_shift_y = abs(curr_car_x - prev_car_x), abs(curr_car_y - prev_car_y)
			pos_shift = np.sqrt(pos_shift_x ** 2 + pos_shift_y ** 2)
			speed = pos_shift / dt

			if speed >= max_speed + 20:
				new_violation_candidates.add(car_id)
				self.num_violations += 1
				print("Warning! Speed limit violation detected.")

				if (len(self.violators_queue) >= 2) and (car_id in self.violators_queue[-1]) and \
					(car_id in self.violators_queue[-2]) and (car_id not in self.last_time_violated or \
						curr_time - self.last_time_violated[car_id] > 60):

					result_violators.append(car_id)
					self.last_time_violated[car_id] = curr_time

		self.violators_queue.append(new_violation_candidates)

		# recognize number plates

		for car_id in result_violators:

			det = self.find_by_id(detection_queue[-1], car_id)
			x_min, y_min, x_max, y_max = det['2d_bbox']
			x_min = max(0, x_min - 50)
			y_min = max(0, y_min - 50)
			x_max = min(frame_queue[-1].shape[1] - 1, x_max + 50)
			y_max = min(frame_queue[-1].shape[0] - 1, y_max + 50)

			crop = frame_queue[-1][x_min: x_max + 1, y_min: y_max + 1, :]
			number_plate = self.number_plate_recognition_service.process(crop)['number_plate']

			if number_plate is not None:
				pickle_data = {'type': "speed limit violation", 'number plate': number_plate, 'frames': frame_queue}
				pickle_path = f"database/speed_limit_violations/{self.num_violations}.pkl"
				pickle.dump(pickle_data, open(pickle_path, 'wb'))


import os

import ember
import lightgbm as lgb  # Gradient Boosting
import shap
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
import pandas as pd
import numpy as np
import pickle
import lief
import json
import torch

# from MalConvModel import MalConvModel
from MalConvModel import MalConvModel
from ember import PEFeatureExtractor
import torch.nn.functional as nn
import tensorflow as tf

def main():
    X_train, y_train, X_test, y_test = ember.read_vectorized_features("ember2018")
    model = lgb.Booster(model_file="ember2018/ember_model_2018.txt")
    model.params['objective'] = 'binary'

    # explainer = shap.TreeExplainer(model, )

    # warnings.simplefilter(action='ignore', category=FutureWarning)
    #
    # def plotImp(model, num=40, fig_size=(40, 20)):
    #     """
    #     Grafica le 40 feature piu' importanti.
    #     """
    #     feature_imp = pd.DataFrame(
    #         {'Value': model.feature_importance(importance_type='gain'), 'Feature': model.feature_name()})
    #     plt.figure(figsize=fig_size)
    #     sns.set(font_scale=2)
    #     sns.barplot(x="Value", y="Feature", data=feature_imp.sort_values(by="Value",
    #                                                                      ascending=False)[0:num])
    #     plt.title('LightGBM Features (avg over folds)')
    #     plt.tight_layout()
    #     plt.savefig('lgbm_importances-01.png')
    #     plt.show()
    #
    # print("> Plot the most 40 important features")
    # plotImp(model)

    # print("> Shap values")
    # shap_values = explainer.shap_values(np.asarray(X_test.tolist())[:100])

    # with open('shap_values', 'wb') as s:
    #     pickle.dump(shap_values, s, protocol=pickle.HIGHEST_PROTOCOL)

    # with open('shap_values', 'rb') as handle:
    #     shap_values = pickle.load(handle)
    #
    # print("> Summary plot shap values")
    # shap.summary_plot(shap_values[0], np.asarray(X_test.tolist())[:100], max_display=20, sort=True)

    data = open('output/AddString_5b_Tequila.exe', "rb").read()
    # data = open('Win32.DarkTequila.exe', "rb").read()
    # data = open('NETFX1.1_bootstrapper.exe', "rb").read()
    extractor = PEFeatureExtractor(2)

    features = np.array(extractor.feature_vector(data), dtype=np.float32)
    with open('output/AddString_5b_Tequila', 'wb') as f:
        pickle.dump(features, f, protocol=pickle.HIGHEST_PROTOCOL)

    # with open('features', 'rb') as handle:
    #     features = pickle.load(handle)

def modify_pe():
    def get_binary(file):
        return lief.parse(file)

    def get_builder(binary):
        return lief.PE.Builder(binary)

    def add_section_constant(binary, name, constant, size):
        # create a section
        section = lief.PE.Section(name)

        # fill it with a constant
        section.content = [constant] * size

        # add the section
        binary.add_section(section, lief.PE.SECTION_TYPES.DATA)

        # build and reparse
        builder = get_builder(binary)
        builder.build()
        builder.write('output/AddConstantTequila.exe')
        binary = get_binary('output/AddConstantTequila.exe')

        return binary

    def add_section_strings(binary, name, string_file):
        # grab strings
        with open(string_file, 'r') as fd:
            data = fd.read()

        # create new section
        section = lief.PE.Section(name)

        # convert characters to decimal representation
        section.content = [ord(c) for c in data]

        # add section to binary
        binary.add_section(section, lief.PE.SECTION_TYPES.DATA)

        # build and reparse
        builder = get_builder(binary)
        builder.build()
        builder.write('output/AddString_5b_Tequila.exe')
        binary = get_binary('output/AddString_5b_Tequila.exe')

        return binary

    def modify_binary(binary):
        header = binary.header
        print(header)
        # print(binary.optional_header)
        print(binary.debug)

    model = lgb.Booster(model_file="ember2018/ember_model_2018.txt")

    # parse internet explorer
    tequila_bin = get_binary('Win32.DarkTequila.exe')

    # add a constant section isn't good => 0.9998018082502094
    # add_section_constant(tequila_bin, "new_section", 250, 500000)

    # 1 meh => 0.8737962293318766
    # add_section_strings(tequila_bin, "license", 'netf.txt')
    # 2 much better => 0.7946982046366664
    # add_section_strings(tequila_bin, "license", 'winrar.txt')
    # 3 very very much better => 0.3025394133707842
    # add_section_strings(tequila_bin, "license", 'notepad.txt')
    # 4 meh => 0.7946982046366664
    # add_section_strings(tequila_bin, "license", 'netf+winrar.txt')
    # 5 OMG => 0.22589567557279205
    add_section_strings(tequila_bin, "license", 'all.txt')

def get_predict_models(name, bytez):
    # thresholds are set here
    print('--------------------')
    print(f'PE name: {name}')
    malconv = MalConvModel('malconv/malconv.checkpoint', thresh=0.5)
    malconv_score = malconv.predict(bytez)
    print(f'{malconv.__name__}: {malconv_score} {malconv_score > malconv.thresh}')

    model = lgb.Booster(model_file="ember2018/ember_model_2018.txt")
    model.params['objective'] = 'binary'
    # Ember  thresh > 0.8336
    ember_score = ember.predict_sample(model, bytez)
    print(f'ember : {ember_score} {ember_score > 0.8336}')
    print('--------------------')

def example():
    WinRAR = open('WinRAR-x64-602it.exe', "rb").read()
    # DarkTequila = open('Win32.DarkTequila.exe', "rb").read()
    # print('File size:', os.path.getsize('Win32.DarkTequila.exe'))
    # AddConstantTequila = open('output/AddConstantTequila.exe', "rb").read()
    # AddString_1_Tequila = open('output/AddString_1_Tequila.exe', "rb").read()
    # AddString_2_Tequila = open('output/AddString_2_Tequila.exe', "rb").read()
    # AddString_3_Tequila = open('output/AddString_3_Tequila.exe', "rb").read()
    # AddString_4_Tequila = open('output/AddString_4_Tequila.exe', "rb").read()
    # AddString_5_Tequila = open('output/AddString_5_Tequila.exe', "rb").read()
    new_DarkTequila = open('mod_5.exe', "rb").read()

    # get_predict_models('WinRAR', WinRAR)
    # get_predict_models('DarkTequila', DarkTequila)

    # get_predict_models('AddConstantTequila', AddConstantTequila)
    # get_predict_models('AddString_1_Tequila', AddString_1_Tequila)
    # get_predict_models('AddString_2_Tequila', AddString_2_Tequila)
    # get_predict_models('AddString_3_Tequila', AddString_3_Tequila)
    # get_predict_models('AddString_4_Tequila', AddString_4_Tequila)
    # get_predict_models('AddString_5_Tequila', AddString_5_Tequila)
    get_predict_models('new_DarkTequila', new_DarkTequila)

if __name__ == '__main__':
    # main()
    # modify_PE()
    example()


import numpy as np
import yaml

import torch
import torch.nn.functional as F
from torchvision.models.video import mc3_18, r2plus1d_18
from facenet_pytorch import MTCNN

from cnn3d import *
from cnn2d import *
from constants import *
from face_utils import *

__all__ = ['Ensemble']

def inference_3d(models_3d, faces, video):
    predictions = []
    try:
        for modeldict in models_3d:
            preds_video = []
            model = modeldict['model']

            for i_face, face in enumerate(faces):
                (frame_from, frame_to), (row_from,
                                         row_to), (col_from, col_to) = face

                x = video[frame_from:frame_to,
                          row_from:row_to + 1, col_from:col_to + 1]
                x = resize_and_square_face(x, output_size=OUTPUT_FACE_SIZE)

                if PRE_INFERENCE_CROP and PRE_INFERENCE_CROP != OUTPUT_FACE_SIZE:
                    x = center_crop_video(x, PRE_INFERENCE_CROP)

                with torch.no_grad():

                    if modeldict['norm'] == '112_imagenet':
                        x = np.array([test_transforms_114_imagenet(
                            image=frame)['image'] for frame in x])
                    elif modeldict['norm'] == '224_imagenet':
                        x = np.array([test_transforms_224_imagenet(
                            image=frame)['image'] for frame in x])

                    x = torch.from_numpy(x.transpose([3, 0, 1, 2])).float()

                    if modeldict['norm'] == 'i3d':
                        x = (x / 255.) * 2 - 1
                    elif modeldict['norm'] == 'nil':
                        pass
                    elif modeldict['norm'] == '112_imagenet':
                        pass
                    elif modeldict['norm'] == '224_imagenet':
                        pass
                    else:
                        raise ValueError(
                            f"Unknown normalisation mode {modeldict['norm']}")

                    y_pred = model(x.cuda())
                    prob0, prob1 = torch.mean(
                        torch.exp(F.log_softmax(y_pred, dim=1)), dim=0)
                    if REVERSE_PROBS:
                        prob1 = 1-prob1
                    preds_video.append(float(prob1))

            if len(preds_video) > 0:
                predictions.append(USE_FACE_FUNCTION(preds_video) * RATIO_3D)
    except:
        pass
    return predictions

def inference_2d(model_2d, faces, coords, video, loader):
    try:
        FRAMES2D = 32
        # Ian's 2D model
    #     coords = coords_by_videopath[videopath]
        preds_video = []

        for i_coord, coordinate in enumerate(coords):
            (frame_from, frame_to), (row_from,
                                     row_to), (col_from, col_to) = faces[i_coord]
            x = []
            for coord_ind, frame_number in enumerate(range(frame_from, min(frame_from+FRAMES2D, frame_to-1))):
                if coord_ind >= len(coordinate):
                    break
                x1, y1, x2, y2 = coordinate[coord_ind]
                x.append(video[frame_number, y1:y2, x1:x2])
            x = np.asarray(x)
            # Reverse back to BGR because it will get reversed to RGB when preprocessed
            # x = x[...,::-1]
            # Preprocess
            x = loader.dataset.process_video(x)
            # x = np.asarray([loader.dataset.process_image(_) for _ in x])
            # Flip every other frame
            x[:, ::2] = x[:, ::2, :, ::-1]
            # RGB reverse every 3rd frame
            # x[:,::3] = x[::-1,::3]
            with torch.no_grad():
                out = model_2d(torch.from_numpy(
                    np.ascontiguousarray(x)).unsqueeze(0).cuda())
            # out = np.median(out.cpu().numpy())
            preds_video.append(out.cpu().numpy())
        if len(preds_video) > 0:
            return USE_FACE_FUNCTION(preds_video) * RATIO_2D
        else:
            pass
    except:
        pass

class Ensemble():
    def __init__(self):
        self.mtcnn = MTCNN(margin=0, keep_all=True, post_process=False, select_largest=False,
                           device='cuda:0', thresholds=MTCNN_THRESHOLDS, factor=MMTNN_FACTOR)
        self.models_3d = build_models()
        with open('./cnn2d/experiment001.yaml') as f:
            CFG = yaml.load(f, Loader=yaml.FullLoader)

            CFG['model']['params']['pretrained'] = None
            model2d = build_model(CFG['model']['name'], CFG['model']['params'])
            model2d.load_state_dict(torch.load(
                './weights/SRXT50_094_VM-0.2504.PTH'))
            model2d = model2d.eval().cuda()
            self.model_2d = model2d
            self.loader = build_dataloader(
                CFG, data_info={'vidfiles': [], 'labels': []}, mode='predict')

    def inference(self, video_path):
        faces, coords = face_detection_wrapper(self.mtcnn, video_path, every_n_frames=FACE_FRAMES,
                                               facedetection_downsample=FACEDETECTION_DOWNSAMPLE,
                                               max_frames_to_load=MAX_FRAMES_TO_LOAD)
        if len(faces):
            last_frame_needed = get_last_frame_needed_across_faces(faces)
            video, rescale = load_video(video_path, every_n_frames=1, to_rgb=True, rescale=None,
                                        inc_pil=False, max_frames=last_frame_needed)
        else:
            return 0.5
        predictions = inference_3d(self.models_3d, faces, video)
        preds_2d = inference_2d(self.model_2d, faces,
                                coords, video, self.loader)
        if preds_2d:
            predictions.append(preds_2d)
        if len(predictions) > 0:
            return np.clip(np.mean(predictions), PROB_MIN, PROB_MAX)
        else:
            return 0.5


import re
import subprocess as sp

import numpy as np
import torch
from more_itertools import chunked

from .bg import DEVICE, Net, iter_frames, remove_many

def basic_greenscreen(path, gpubatchsize, model_name, frame_limit=-1):
    command = None
    proc = None
    net = Net(model_name)
    script_net = None
    for gpu_batch in chunked(iter_frames(path), gpubatchsize):
        if 0 <= frame_limit < gpubatchsize:
            break
        frame_limit -= gpubatchsize
        if script_net is None:
            script_net = torch.jit.trace(net, torch.as_tensor(np.stack(gpu_batch), dtype=torch.float32, device=DEVICE))
        for image in remove_many(gpu_batch, script_net):
            if command is None:
                command = ['FFMPEG',
                           '-y',
                           '-f', 'rawvideo',
                           '-vcodec', 'rawvideo',
                           '-s', F"{image.shape[1]}x320",
                           '-pix_fmt', 'gray',
                           '-r', "30",  # for now I am hardcoding it, I can always resize the clip in premiere anyway
                           '-i', '-',
                           '-an',
                           '-vcodec', 'mpeg4',
                           '-b:v', '2000k',
                           re.sub(r"\.(mp4|mov|avi)", r".matte.\1", path, flags=re.I)]
                proc = sp.Popen(command, stdin=sp.PIPE)
            proc.stdin.write(image.tostring())
    if isinstance(proc, sp.Popen):
        proc.stdin.close()
        proc.wait()


import argparse
import json
import os

import librosa
import matplotlib.pyplot as plt
import numpy as np
import torch

from config import Config
from retriever import Retriever
from utils.hifigan import HiFiGANWrapper
from utils.libritts import DumpedLibriTTS
from speechset.utils.melstft import MelSTFT

parser = argparse.ArgumentParser()
parser.add_argument('--ret-config')
parser.add_argument('--ret-ckpt')
parser.add_argument('--hifi-config')
parser.add_argument('--hifi-ckpt')
parser.add_argument('--data-dir')
parser.add_argument('--audio1')
parser.add_argument('--audio2')
args = parser.parse_args()

with open(args.ret_config) as f:
    config = Config.load(json.load(f))

ckpt = torch.load(args.ret_ckpt, map_location='cpu')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

retriever = Retriever(config.model)
retriever.load(ckpt)
retriever.to(device)
retriever.eval()

hifigan = HiFiGANWrapper(args.hifi_config, args.hifi_ckpt, device)

if args.data_dir is not None:
    libritts = DumpedLibriTTS(args.data_dir)
    testset = libritts.split(config.train.split)

    START = 400
    sid, text, mel, textlen, mellen = testset[START:START + 2]
else:
    def mel_fn(path: str, stft: MelSTFT = MelSTFT(config.data)) -> np.ndarray:
        wav, _ = librosa.load(path)
        return stft(wav)
    # pack
    mel = [mel_fn(path) for path in [args.audio1, args.audio2]]
    mellen = np.array([m.shape[0] for m in mel], dtype=np.int64)
    mel = np.stack([np.pad(m, [[0, mellen.max() - m.shape[0]], [0, 0]]) for m in mel])

# wrap
mel, mellen = torch.tensor(mel, device=device), torch.tensor(mellen, device=device)
# clip
MAX_MELLEN = 240
mel, mellen = mel[:, :MAX_MELLEN], torch.clamp_max(mellen, MAX_MELLEN)

with torch.no_grad():
    synth, aux = retriever(mel, mellen)
    # flip style
    style = aux['style'][[1, 0]]

    mixed, _ = retriever(mel, mellen, refstyle=style)
    # vocoding
    out = hifigan.forward(torch.cat([synth, mixed], dim=0))
    out = out.cpu().numpy()

    # [B, G, V, T]
    indices = torch.zeros_like(aux['logits']).scatter(
        2, aux['logits'].argmax(dim=2, keepdim=True), 1.).cpu().numpy()

os.makedirs('./synth', exist_ok=True)
for i, (wav, mlen) in enumerate(zip(out, mellen.repeat(2))):
    # unwrap
    librosa.output.write_wav(
        f'./synth/{i}.wav', wav[:mlen.item() * config.data.hop], sr=config.data.sr)

cmap = np.array(plt.get_cmap('viridis').colors)
# B x ([G, V, T], [T x R, mel])
for i, (idx, mel) in enumerate(zip(indices, synth.cpu().numpy())):
    rest = mel.shape[0] % config.model.reduction
    if rest > 0:
        mel = np.pad(mel, [[0, config.model.reduction - rest], [0, 0]], constant_values=np.log(config.data.eps))
    # [T, mel]
    mel = mel.reshape(-1, config.model.reduction, config.data.mel).mean(axis=1)
    # [T, mel]
    mel = (mel - mel.min()) / (mel.max() - mel.min() + 1e-7)
    # [mel, T]
    mel = np.flip(mel.transpose(1, 0), axis=0)
    # G x [V, T]
    for j, imap in enumerate(idx):
        # [mel + V, T]
        img = np.concatenate([mel, imap], axis=0)
        # [mel + V, T, 3]
        plt.imsave(f'./synth/{i}_{j}.png', cmap[(img * 255).astype(np.uint8)])


import numpy as np

class TrustFile(object):
    _COMPO_NAMES = ["X", "Y", "Z","4","5","6","7","8","9"]
    """ Parsing and value extraction from a Trust file. Assumption: all valid (=completly written) line
    of data terminate with a new line char """

    _VOID_ARR = np.array([], dtype=np.float32)

    def __init__(self, filePath, refFilePath):
        """
        @param filePath Trust file to parse
        @param refFilePath file serving as a reference to know if the main file has been reset. This
        is typically the stop file. Can be None for test purposes (in this case only size of the main file
        will be inspected).
        """
        self._filePath = filePath
        self._refPath = refFilePath
        self._lastSize = -1
        self._lastRefDate = -1
        self._reset()

    def _reset(self):
        self._entries = {}                # key = entry_name, val=col_index (or component numb for segs)
        self._taxis = TrustFile._VOID_ARR # time axis values
        self._yvalues = TrustFile._VOID_ARR           # a np.array, col index is in _entries
        self._lastIndex = {}              # position of the last fetched data in the yvalues arrays. key=entry_name, vals=index
        self._lastSeek = 0                # position of the last fetched data in the file
        self._data = ""                   # last chunk of read data
        self._ncols = 0                   # numb of cols in the data

    def isLogY(self):
        return False

    def getXLabel(self):
        raise NotImplementedError

    def getEntries(self):
        """
        @return a list of file entries (i.e. what can be ploted from the file)
        """
        status = self._queryAndUpdateStatus(update=False)
        if status == "invalid":
            return []
        if not len(self._entries):
            hdr = self._populateFromHeader()
            self._reloadEntries(hdr)
        return list(self._entries.keys())

    def getValues(self, entryName):
        """
        Get all current valid values from the file.
        @return two np arrays (x and y)
        """
        fileReady = self._getValuesInternal()
        return self._extractValues(fileReady, entryName)

    def getNewValues(self, entryName):
        """
        Return values to be appended to the last retrieved ones to get the full set of data.
        Note: after a file reset, it will return None, None on an entryName as long as getValues() hasn't been called once.
        @return two (potentially void) np arrays (x and y)
        @return None, None if the file has been reset or has become invalid (getValues() must be invoked again on this entry)
        """
        fileReady = self._getValuesInternal()
        return self._extractNewValues(fileReady, entryName)

    def _getValuesInternal(self):
        """ @return False if file is not ready """
        status = self._queryAndUpdateStatus(update=True)
#    print "getValues stat", status, " ", self._filePath
        if status == "invalid":
            return False
        if not len(self._entries):   # if there's been a reset()
            hdr = self._populateFromHeader()
            self._reloadEntries(hdr=hdr)
        # Try to retreive values only if entries are OK:
        if not len(self._entries):
            # File is not ready:
            return False
        if not self._yvalues.shape[0]:
            self._reloadValues()
        elif status != "same":
            self._completeValues()
        if not self._yvalues.shape[0]:
            # File is not ready:
            return False
        return True

    def _extractValues(self, fileReady, entryName):
        """ Extract values from the internal array """
        if not fileReady:
            # Two void arrays
            return TrustFile._VOID_ARR, TrustFile._VOID_ARR
        else:
            idx = self._taxis.shape[0]
            self._lastIndex[entryName] = idx
            col = self._entries[entryName]
            return self._taxis, self._yvalues[:, col]

    def _extractNewValues(self, fileReady, entryName):
        """ Extract NEW values from the internal array """
        if not fileReady:
            return None, None
        else:
            idx = self._lastIndex[entryName]
            if idx == -1:
                # File was reset previously, and for this entry, getValues() has not been called again yet
                return None, None
            col = self._entries[entryName]
            x, y = self._taxis[idx:], self._yvalues[idx:, col]
            self._lastIndex[entryName] += x.shape[0]
            return x, y

    def _populateFromHeader(self, hdr=None):
        """ Fill meta data from file header. If hdr=None, the method should reread the header itself.
        @return hdr if it could be read properly, None otherwise
        """
        return hdr

    def _reloadEntries(self, hdr=None):
        """ Reload entries exposed by a file
        @param hdr if None the header has to be re-read by the method
        """
        pass

    def _reloadValues(self):
        """ Reload all data (not header) from a file.
        Use numpy.loadtxt for efficiency. """
        self._yvalues = self._readOn()
        if self._yvalues.shape[0]:
            self._taxis = self._yvalues[:,0]
        else:
            self._taxis = TrustFile._VOID_ARR

    def _completeValues(self):
        """ Complete the internal array with the newly read values """
        lines = self._readOn()
        if lines.shape[0]:
            self._taxis = np.append(self._taxis, lines[:,0])
            if self._yvalues.shape[0]:
                self._yvalues = np.vstack([self._yvalues, lines])
            else:
                self._yvalues = lines

    def _queryAndUpdateStatus(self, update=True):
        """
        @param update True if the internal status of the file should be updated
        @return "append" : if new data has been added (or if first query on a valid file)
                    "reset"  : if file has been reset
                    "invalid": if the file cannot be accessed anymore or is invalid
                    "same"   : if the file has not changed
        """
        import os
        try:
            ret = "invalid"
            if self._refPath is not None:
                t = os.stat(self._refPath).st_mtime
            else:
                # for test - will emulate a single reset at startup
                t = -2
            if t != self._lastRefDate:
                ret = "reset"
                if update:
                    self._lastRefDate = t
                    self._reset()
            siz = os.stat(self._filePath).st_size
            if ret == "reset":
                if update:     self._lastSize = siz
                return ret
            if siz < self._lastSize:
                if update:
                    self._lastSize = siz
                    self._reset()
                print("WARNING: file is smaller but reference file has not changed! Will reset. (%s)" % self._filePath)
                return "reset"
            elif siz == self._lastSize:
                return "same"
            else:
                if update:
                    self._lastSize = siz
                return "append"
        except OSError:
            if update:
                self._lastRefDate = -1
                self._lastSize = -1
                self._reset()
            return "invalid"

    def _parseData(self, data):
        arrPy = data.split("\n")
        remainData, npArr = "", TrustFile._VOID_ARR
        if len(arrPy):
            remainData = arrPy[-1]
        else:
            return remainData, npArr
        i = 0
        # Skip header:
        for a in arrPy[:-1]:
            if a[0] == "#":
                i += 1
            else:
                break
        # No data, just a header:
        if len(arrPy[i:-1]) == 0:
            return remainData, npArr
        arrData = " ".join(arrPy[i:-1])
        arrData = np.fromstring(arrData, dtype=np.float32, sep=' ')
        if arrData.shape[0] % self._ncols:
            raise Exception("Misformatted file: the line below has an invalid numb of cols:\n%s" % a)
        npArr = arrData.reshape((-1, self._ncols))
        return remainData, npArr

    def _readOn(self):
        """ Suppose the file can be read (_queryAndUpdateStatus has been invoked)
        @return a list of NumPy arrays.
        """
        try:
            with open(self._filePath, "r") as f:
                f.seek(self._lastSeek)
                data = f.read()
                self._lastSeek += len(data)
                # Now try to extract lines:
                self._data += data
                self._data, lines = self._parseData(self._data)
                return lines
        except IOError:
            return None

    @classmethod
    def _ReadHeader(cls, filePth, nLines):
        """ Read header (nLines-1 first lines) + one data line (to know numb of compo.) """
        ret = []
        try:
            with open(filePth) as f:
                for _ in range(nLines):
                    l = f.readline().strip()
                    if l.startswith("#"):
                        l = l[1:].strip()
                    if l != "":
                        ret.append(l.split())
        except:
            return None
        return ret

class SonFile(TrustFile):

    def __init__(self, filePath, refPath):
        TrustFile.__init__(self, filePath, refPath)
        self._reset()

    def _reset(self):
        TrustFile._reset(self)
        self._dim = -1                       # Space dimension of the points
        self._ncompo = -1                    # Numb of components of the field
        self._points = TrustFile._VOID_ARR   # The points themselves
        self._field = ""    # Field name (VITESSE, etc ...)

    @classmethod
    def BuildFromHeader(cls, filePath, refPath):
        """ Build either a SonPOINTFile or SonSEGFile depending on the header
        """
        try:
            hdr = TrustFile._ReadHeader(filePath,5)
            if hdr is None:
                return None
            typ = hdr[3][1]
            if typ == "POINTS" or typ == "NUMERO_ELEM_SUR_MAITRE":
                ret = SonPOINTFile(filePath, refPath)
            elif typ == "SEGMENT":
                ret = SonSEGFile(filePath, refPath)
            else:
                raise Exception("Invalid SON file type!")
        except:
            return None, None
        ret._populateFromHeader(hdr)
        return ret, hdr

    def _populateFromHeader(self, hdr=None):
        if hdr is None:
            hdr = TrustFile._ReadHeader(self._filePath,5)
        try:
            # Field name:
            self._field = hdr[2][1]
            # Points:
            dim = 1
            if len(hdr[1]) > 2:
                if hdr[1][3] == "y=": dim += 1
                if len(hdr[1]) > 5:
                    if hdr[1][5] == "z=": dim += 1
            self._dim = dim
            # Store point coords:
            if (len(hdr[1])-1) % 2 or ((len(hdr[1])-1)//2) % dim:
                return None
            npoints = (len(hdr[1])-1)//2//dim
            pts = []
            for i in range(npoints):
                pt = [float(a) for a in hdr[1][2+i*2*dim:2+(i+1)*2*dim:2]]
                pts.append(pt)
            # And finally number of components:
            if (len(hdr[4]) - 1) % npoints:
                return None
            self._ncompo = (len(hdr[4]) - 1) // npoints
            self._points = np.array(pts)
            self._ncols = len(hdr[4])
            return hdr
        except:
            self._reset()
            # A partial/misformatted file is ignored:
            return None

    def _displayPoint(self, pt):
        pts = ["%s=%g" % (SonFile._COMPO_NAMES[i], pt[i]) for i in range(len(pt))]
        return ",".join(pts)

    def _reloadEntries(self, hdr=None):
        if hdr is None:
            hdr = TrustFile._ReadHeader(self._filePath, 5)
        #self._entries[hdr[1][0]] = 0  # do not take time
        col = 1
        for p in self._points:
            sp = self._displayPoint(p)
            for cnum in range(self._ncompo):
                if self._ncompo > 1:
                    e = "%s_%s (%s)" % (self._field, SonFile._COMPO_NAMES[cnum], sp)
                else:
                    e = "%s (%s)" % (self._field, sp)
                self._entries[e] = col
                col += 1
        # Reset last index:
        self._lastIndex = dict( list(zip(list(self._entries.keys()), [-1] * len(self._entries))) )

class SonPOINTFile(SonFile):
    def __init__(self, filePath, refPath):
        SonFile.__init__(self, filePath, refPath)

    def getXLabel(self):
        return "Time"

class SonSEGFile(SonFile):
    """ For a segment file, getNewValues() will return None, None everytime the last line has changed.
     Hence by default the API of a seg file return  """
    def __init__(self,filePath, refPath):
        SonFile.__init__(self, filePath, refPath)
        self._reset()

    def _reset(self):
        SonFile._reset(self)
        self._start = []
        self._end = []
        self._orient = ""                  #  "x", "y", "z"  or "curvi."
        self._xaxis = TrustFile._VOID_ARR  # X axis for seg plots
        self._segentries = {}

    def _populateFromHeader(self, hdr=None):
        hdr = SonFile._populateFromHeader(self, hdr)
        if not hdr is None:
            l = [0.0] * (2 * self._dim)
            l2 = [float(hdr[3][i]) for i in range(2,2+2*self._dim)]
            start, end = l2[:self._dim], l2[self._dim:]
            self.setXTremePoints(start, end)
            self._computeXAxis()
            return hdr
        else:
            return None

    def _computeCurvilinear(self):
        """ Compute curvilinear distances """
        d = self._points[1:] - self._points[:-1]
        dist = np.sqrt(np.sum(d*d, axis=1))
        curvi = np.zeros((self._points.shape[0],), dtype=np.float32)
        curvi[1:] = np.cumsum(dist)
        return curvi

    def _computeXAxis(self):
        if self._orient == "curvi.":
            self._xaxis = self._computeCurvilinear()
        else:
            mp = {"x": 0, "y": 1, "z": 2}
            self._xaxis = self._points[:, mp[self._orient]]

    def setXTremePoints(self, start, end):
        x1, x2 = start[0], end[0]
        if self._dim > 1:   y1, y2 = start[1], end[1]
        else:               y1, y2 = 0.0, 0.0
        if self._dim > 2:   z1, z2 = start[2], end[2]
        else:               z1, z2 = 0.0, 0.0
        if x1 == x2 and y1 == y2:
            self._orient = "z"
        elif x1 == x2 and z1 == z2:
            self._orient = "y"
        elif y1 == y2 and z1 == z2:
            self._orient = "x"
        else:
            self._orient = "curvi."
        self._start, self._end = start, end

    def getXLabel(self):
        return "Coordinate (%s)" % self._orient

    def _extractValues(self, fileReady, entryName):
        """ Override to handle seg entries. """
        if entryName in self._entries:
            return SonFile._extractValues(self, fileReady, entryName)
        elif entryName in self._segentries:
            return self._extractValuesSeg(fileReady, entryName)
        else:
            raise Exception("Invalid column name! '%s'" % entryName)

    def _extractNewValues(self, fileReady, entryName):
        """ Override to handle seg entries - the behaviour is also to return None, None everytime a new seg is available. """
        if entryName in self._entries:
            return SonFile._extractNewValues(self, fileReady, entryName)
        elif entryName in self._segentries:
            return self._extractNewValuesSeg(fileReady, entryName)
        else:
            raise Exception("Invalid column name! '%s'" % entryName)

    def _extractValuesSeg(self, fileReady, entryName):
        if not fileReady:
            return TrustFile._VOID_ARR, TrustFile._VOID_ARR
        else:
            col = self._segentries[entryName]
            self._lastIndex[entryName] = 1  # 1, no more new data available
            x, y = self._xaxis, self._yvalues[-1, 1+col::self._ncompo]
            return x, y

    def _extractNewValuesSeg(self, fileReady, entryName):
        if not fileReady:
            return None, None
        else:
            # Segment data - getNewValues()'s behaviour is to indicate with (None, None) that a new seg is available
            alreadyFetched = self._lastIndex[entryName]
            if not alreadyFetched:
                x, y = None, None
            else:
                x, y = TrustFile._VOID_ARR, TrustFile._VOID_ARR   # Nothing has changed
            return x, y

    def getEntriesSeg(self):
        """ Return only seg entries (i.e. not normal point entries) """
        entr = self.getEntries()  # make sure getEntries logic is applied
        return list(self._segentries.keys())

    def getPointEntries(self, segName):
        self.getEntries()  # make sure getEntries logic is applied
        compoNum = self._segentries[segName]
        l = []
        for e in self._entries:
            if self._ncompo == 1 or (self._entries[e] - compoNum) % self._ncompo:
                l.append(e)
        return l

    def getLastTime(self):
        if self._yvalues.shape[0]:
            return self._yvalues[-1,0]
        else:
            return 0.0

    def _reloadEntries(self, hdr=None):
        """ Override - append seg entries which are stored with neg indices """
        SonFile._reloadEntries(self, hdr)
        start = self._displayPoint(self._start)
        end = self._displayPoint(self._end)
        for cnum in range(self._ncompo):
            if self._ncompo == 1:
                e = "%s (SEG [%s] -> [%s])" % (self._field, start, end)
            else:
                e = "%s_%s (SEG [%s] -> [%s])" % (self._field, SonFile._COMPO_NAMES[cnum], start, end)
            self._segentries[e] = cnum
            self._lastIndex[e] = 0     # 0, new data available, not fetched yet

    def _completeValues(self):
        SonFile._completeValues(self)
        for s in list(self._segentries.keys()):
            self._lastIndex[s] = 0     # 0, new data available, not fetched yet

class OutFile(TrustFile):

    def __init__(self, filePath, refPath):
        TrustFile.__init__(self, filePath, refPath)
        self._reset()

    def _reset(self):
        TrustFile._reset(self)
        self._dim = -1                       # Space dimension of the points
        self._ncompo = -1                    # Numb of components of the field
        self._field = ""    # Field name (VITESSE, etc ...)

    def getXLabel(self):
        return "Time"
    def _populateFromHeader(self, hdr=None):
        if hdr is None:
            hdr = TrustFile._ReadHeader(self._filePath,4)
        try:

            # print "uuuuuuuuuu",hdr
            # Field name:
            self._fied="Unknown"
            for x in  hdr[1]:
                if x.find("Integral")> -1:
                    self._field = x.split("=")[-1]
            nb_co=len(hdr[2])-1
            # print "uu",nb_co
            dim=1
            lg=len(hdr[3])
            if (lg >1):
                # dim=(len(hdr[3])-1)/nb_co soucis avec perio
                if (hdr[3][3]=='Fx'):
                    dim=2
                else:
                    assert(hdr[3][3]=='Fz')
                    dim=3

            #if (self._field=="Integral(P*ndS)"):
            #  print "pb fichier pression"
            #  nb_co-=1
            # nb_co-=1 # on a total par erreur dans les label
            # print "yyyyyyyyyy",(len(hdr[3])-1)/dim, nb_co
            # print "dim",dim,nb_co
            if dim==0:
                dim=1
            self._dim=dim
            self._ncols = dim*(nb_co)+1
            # self._ncols = len(hdr[4])

            #self._entries["Time"]=0
            col=1
            xold=None
            for c in range(nb_co):
                x=hdr[2][c+1]
                # dans le cas periodique le nom du bord est double
                if x==xold:
                    x+="_bis"
                xold=x
                if dim==1:
                    self._entries[self._field+" "+x]=col
                    col+=1
                else:
                    for i in range(self._dim):
                        # print"uuu", i,x
                        self._entries[self._field+" "+x+OutFile._COMPO_NAMES[i]]=col
                        col+=1
            # print "uuu",self._entries
            return hdr
        except:
            self._reset()
            # A partial/misformatted file is ignored:
            return None

    def _reloadEntries(self, hdr=None):
        if hdr is None:
            self._populateFromHeader(hdr)
            # print "uuuuuuuuuuuu"
            #1/0

        self._lastIndex = dict( list(zip(list(self._entries.keys()), [-1] * len(self._entries))) )

class CSVFile(TrustFile):
    def __init__(self,filePath, nbclos,refPath=None):
        TrustFile.__init__(self, filePath, refPath)
        self._ncols = nbclos
        self._ncols_sa = nbclos
    def reset(self):
        print("reset")
        pass
    def getXLabel(self):
        return "??"

    def _populateFromHeader(self, hdr=None):
        self._ncols =  self._ncols_sa
        if (hdr is None):
            h1=[]
            for i in range(self._ncols):
                h1.append("c%d"%i)
            hdr=[h1]

        for i, col in enumerate(hdr[0]):
            self._entries[col] = i
        return hdr

    def _reloadEntries(self, hdr=None):

        if hdr is None:
            print("euh...;")
            1/0
        # Adrien ???
        self._lastIndex = dict( list(zip(list(self._entries.keys()), [-1] * len(self._entries))) )
        # print "QQ",self._lastIndex
        # print "uuu",self._entries.keys()

class DTEVFile(TrustFile):
    def __init__(self,filePath, refPath):
        TrustFile.__init__(self, filePath, refPath)

    def isLogY(self):
        return True

    def getXLabel(self):
        return "Time"

    def _populateFromHeader(self, hdr=None):
        if hdr is None:
            hdr = TrustFile._ReadHeader(self._filePath,2)
        self._ncols = len(hdr[0])
        return hdr

    def _reloadEntries(self, hdr=None):
        if len(hdr[0]) >= 6:
            for i, col in enumerate(hdr[0]):
                self._entries[col] = i
        self._lastIndex = dict( list(zip(list(self._entries.keys()), [-1] * len(self._entries))) )

def BuildFromPath(pth, ref_path):
    """ Object factory. """
    if pth[-4:] == ".son":
        ret, hdr = SonFile.BuildFromHeader(pth, ref_path)
        return ret
    elif pth[-6:] == ".dt_ev":
        return DTEVFile(pth, ref_path)
    elif pth[-4:] == ".out":
        return OutFile(pth, ref_path)

    else:
        raise ValueError("Invalid result file type!")


import math
import numpy
import random
from PIL import Image, ImageChops
from util import jpeg_header_length, jpeg_image_from_file, random_block
from constants import BLOCK_SIZE

################################
# IMAGE DEGRADING METHODS
################################

LAST_SEED_USED = 0

def degrade_jpeg_byte_array(byte_array):
    """Sets some of the bytes in the given jpeg byte array to zero.
    Number of bytes affected depends on the size of the array."""
    # Try not to edit the jpeg header.
    header_length = jpeg_header_length(byte_array)
    max_index = len(byte_array) - header_length - 4
    iterations = max(10, int(len(byte_array) / 50000))
    seed = random.randint(0, 100)
    global LAST_SEED_USED
    while seed == LAST_SEED_USED:
        seed = random.randint(0, 100)
    LAST_SEED_USED = seed
    seed_percent = seed / 100
    for i in range(iterations):
        min_pixel_index = math.floor(max_index / iterations * i)
        max_pixel_index = math.floor(max_index / iterations * (i + 1))
        delta = max_pixel_index - min_pixel_index
        pixel_index = math.floor(min_pixel_index + delta * seed_percent)
        if pixel_index > max_index:
            pixel_index = max_index
        index_in_byte_array = math.floor(header_length + pixel_index)
        byte_array[index_in_byte_array] = 0x00
    return byte_array

def degrade_jpeg(input_file, output_file):
    """Randomly edits bytes in the given jpeg file and writes the result to the output file."""
    with open(input_file, 'rb') as f:
        input_byte_array = bytearray(f.read())
    output_byte_array = degrade_jpeg_byte_array(input_byte_array)
    with open(output_file, 'wb') as f:
        f.write(output_byte_array)

def fake_degrade_jpeg(input_file, output_file, strong = False):
    """Manipulates the image to achieve a similar effect to degrade_jpeg, but without
    corrupting the file."""
    image = jpeg_image_from_file(input_file)
    iterations = 10
    options = [fake_adjust_color, fake_offset]
    weights = [0.8, 0.2]
    for i in range(iterations):
        option = numpy.random.choice(options, p = weights)
        image = option(image, strong = strong)
    image.save(output_file, quality = 10, optimize = True)

def fake_adjust_color(image, strong = False):
    """Randomly adjusts the RGB values of part of the image."""
    component = random.randint(0, 2)
    lower_bound = 30 if strong else 10
    upper_bound = 80 if strong else 40
    amount = random.randint(lower_bound, upper_bound)
    start_column, start_row = random_block(image)
    array = numpy.zeros((image.height, image.width, 3), dtype = numpy.uint8)
    array[start_row:(start_row + BLOCK_SIZE), start_column:, component].fill(amount)
    if start_row + BLOCK_SIZE < image.height:
        array[(start_row + BLOCK_SIZE):, :, component].fill(amount)
    map_image = Image.fromarray(array, 'RGB')
    # Add or subtract
    if random.choice([True, False]):
        return ImageChops.add(image, map_image)
    else:
        return ImageChops.subtract(image, map_image)

def fake_offset(image, strong = False):
    """Randomly moves parts of the image right or left."""
    start_column, start_row = random_block(image)
    cropped_image = image.crop((0, start_row, image.width, image.height))
    offset_multiple = random.randint(5, 15) if strong else 1
    offset_size = offset_multiple * BLOCK_SIZE
    offset = random.choice([-1, 1]) * offset_size
    offset_image = ImageChops.offset(cropped_image, offset, 0)
    image.paste(offset_image, (0, start_row))
    return image

################################
# TEXT DEGRADING METHODS
################################

def degrade_text(text):
    """Adds several typographic errors to the given string and returns the result.
    The number of errors added depends on the length of the string."""
    if len(text) < 2:
        return ''
    iterations = max(2, int(len(text) / 8))
    options = [delete_character, delete_character_leaving_space, swap_two_characters, duplicate_character, insert_character_from_text]
    weights = [0.3, 0.1, 0.2, 0.2, 0.2] if len(text) >= 8 else [0, 0, 0.4, 0.3, 0.3]
    for _ in range(iterations):
        option = numpy.random.choice(options, p = weights)
        text = option(text)
    return text

def swap_two_characters(text):
    """Swaps two adjacent characters in the given string and returns the result."""
    index = random.randint(0, len(text) - 2)
    return text[:index] + text[index + 1] + text[index] + text[(index + 2):]

def duplicate_character(text):
    """Doubles up a character in the given string and returns the result.
    e.g. "example" -> "examplle"
    """
    index = random.randint(0, len(text) - 1)
    return text[:index] + text[index] + text[index] + text[(index + 1):]

def insert_character_from_text(text):
    """Selects a character in the given string, inserts a copy of that character
    somewhere else in the string, and returns the result."""
    new_character = random.choice(text)
    index = random.randint(0, len(text))
    return text[:index] + new_character + text[index:]

def delete_character(text):
    """Deletes a character from the given string and returns the result."""
    index = random.randint(0, len(text) - 1)
    return text[:index] + text[(index + 1):]

def delete_character_leaving_space(text):
    """Deletes a character from the given string, leaving a space. Returns the result."""
    index = random.randint(0, len(text) - 1)
    return text[:index] + ' ' + text[(index + 1):]


import torch
import pickle
import numpy as np
from rdkit import Chem

from retroformer.utils.smiles_utils import smi_tokenizer, clear_map_number, SmilesGraph, add_mapping
from retroformer.utils.smiles_utils import canonical_smiles, get_reaction_centers_from_template, select_diverse_candidate
from retroformer.utils.translate_utils import batch_infer_reaction_center
from retroformer.models.model import RetroModel
from retroformer.rdchiral.template_extractor import get_fragments_for_changed_atoms

def gaussian(x, mean, amplitude, standard_deviation):
    return amplitude * np.exp(- (x - mean) ** 2 / (2 * standard_deviation ** 2))

class RetroFormer():
    def __init__(self, checkpoint_path=None, vocab_path=None, rc_path=None, device='cpu'):
        self.device = device
        if vocab_path is not None:
            self.load_vocab(vocab_path)
        if checkpoint_path is not None:
            self.load_model(checkpoint_path)
        if rc_path is not None:
            with open(rc_path, 'rb') as f:
                self.rc_candidates = pickle.load(f)

        self.factor_func = lambda x: (1 + gaussian(x, 5.55391565, 0.27170542, 1.20071279))

    def load_model(self, checkpoint_path):
        self.model = RetroModel(num_layers=8, d_model=256,
                                heads=8, d_ff=2048, dropout=0,
                                vocab_size_src=len(self.vocab_src_stoi),
                                vocab_size_tgt=len(self.vocab_tgt_stoi),
                                shared_vocab=True,
                                shared_encoder=False,
                                src_pad_idx=self.vocab_src_stoi['<pad>'],
                                tgt_pad_idx=self.vocab_tgt_stoi['<pad>'])

        self.model.load_state_dict(torch.load(checkpoint_path)['model'])
        self.model.to(self.device)
        self.model.eval()

    def load_vocab(self, vocab_path):
        with open(vocab_path, 'rb') as f:
            self.vocab_src_itos, self.vocab_tgt_itos = pickle.load(f)
        self.vocab_src_stoi = {v: i for i, v in enumerate(self.vocab_src_itos)}
        self.vocab_tgt_stoi = {v: i for i, v in enumerate(self.vocab_tgt_itos)}

    def reconstruct_smi(self, tokens, src=True, raw=False):
        if src:
            if raw:
                return [self.vocab_src_itos[t] for t in tokens]
            else:
                return [self.vocab_src_itos[t] for t in tokens if t != 1]
        else:
            if raw:
                return [self.vocab_tgt_itos[t] for t in tokens]
            else:
                return [self.vocab_tgt_itos[t] for t in tokens if t not in [1, 2, 3]]

    def parse_inputs(self, mol_smiles):
        cano_prod = clear_map_number(mol_smiles)
        smiles_graph = SmilesGraph(cano_prod)
        src_token = ['<UNK>'] + smi_tokenizer(cano_prod)
        src_token = [self.vocab_src_stoi.get(st, self.vocab_src_stoi['<unk>']) for st in src_token]
        return cano_prod, src_token, smiles_graph

    def tag_reaction_center(self, src_tokens, rc):
        for i in range(len(src_tokens)):
            if i in rc:
                if src_tokens[i][0] == '[':
                    src_tokens[i] = src_tokens[i].replace(']', ':{}]'.format(999))
                else:
                    src_tokens[i] = '[{}:{}]'.format(src_tokens[i], 999)
        smiles = ''.join(src_tokens)
        return smiles

    def infer_reaction_center(self, atom_rc_scores, bond_rc_scores, graph_packs,
                              beta, percent_aa, percent_ab, k=5, num_removal=5, max_count=25):
        _, predicts = batch_infer_reaction_center(atom_rc_scores, bond_rc_scores, graph_packs,
                                                  alpha_atom=-1, alpha_bond=-1, beta=beta,
                                                  percent_aa=percent_aa, percent_ab=percent_ab,
                                                  verbose=False, k=k, factor_func=self.factor_func,
                                                  num_removal=num_removal, max_count=max_count)
        return predicts

    def build_rc_smarts(self, cano_prod, non_reactive_mask):
        cano_prod_explicitH = Chem.MolToSmiles(Chem.MolFromSmiles(cano_prod), allHsExplicit=True)
        prod_tokens = smi_tokenizer(cano_prod_explicitH)
        for i, token in enumerate(prod_tokens):
            if not non_reactive_mask[i]:
                if token[0] == '[' and token[-1] == ']':
                    prod_tokens[i] = prod_tokens[i].replace(']', ':{}]'.format(i))
                else:
                    prod_tokens[i] = '[{}:{}]'.format(prod_tokens[i], i)

        mol = Chem.MolFromSmiles(''.join(prod_tokens))

        atom_tags = []
        for atom in mol.GetAtoms():
            if atom.HasProp('molAtomMapNumber'):
                atom_tags.append(atom.GetProp('molAtomMapNumber'))

        prod_rc_smarts, _, _ = get_fragments_for_changed_atoms([mol], atom_tags, category='product')
        prod_rc_smarts = prod_rc_smarts[1:-1]

        patt = Chem.MolFromSmarts(prod_rc_smarts)
        for atom in patt.GetAtoms():
            if atom.HasProp('molAtomMapNumber'):
                atom.ClearProp('molAtomMapNumber')

        rc_smarts = Chem.MolToSmarts(patt)
        #
        # except:
        #     atom_to_use, atom_symbols = [], []
        #     for atom in mol.GetAtoms():
        #         if atom.HasProp('molAtomMapNumber'):
        #             atom_to_use.append(atom.GetIdx())
        #             atom.ClearProp('molAtomMapNumber')
        #         atom_symbols.append(atom.GetSmarts())
        #
        #     rc_smarts = Chem.MolFragmentToSmiles(mol, atom_to_use,
        #                                          atomSymbols=atom_symbols, allHsExplicit=True,
        #                                          isomericSmiles=True, allBondsExplicit=True)

        return rc_smarts

    def compute(self, product_smiles, verbose=False, topk=5, use_template=False):
        cano_prod, src_token, smiles_graph = self.parse_inputs(product_smiles)
        # Prepare Inputs:
        src = torch.LongTensor(src_token).unsqueeze(1).to(self.device)
        full_adj_matrix = torch.from_numpy(smiles_graph.full_adjacency_tensor)
        bond = torch.zeros((1, len(src_token), len(src_token), 7), dtype=torch.long).to(self.device)
        bond[0, 1:full_adj_matrix.shape[0] + 1, 1:full_adj_matrix.shape[1] + 1] = full_adj_matrix
        graph_packs = (bond, [smiles_graph])

        # Encode:
        with torch.no_grad():
            prior_encoder_out, edge_feature = self.model.encoder(src, bond)
            atom_rc_scores = self.model.atom_rc_identifier[1](self.model.atom_rc_identifier[0](prior_encoder_out) / 10)
            bond_rc_scores = self.model.bond_rc_identifier[1](self.model.bond_rc_identifier[0](edge_feature) / 10)

        if use_template:
            pair_indices = torch.where(bond.sum(-1) > 0)
            batch_bond_scores = torch.zeros(src.shape[1], src.shape[0], src.shape[0]).to(src.device)
            batch_bond_scores[pair_indices] = bond_rc_scores.view(-1)
            batch_atom_scores = atom_rc_scores.squeeze(2).transpose(0, 1)

            src_tokens = self.reconstruct_smi(src[:, 0], src=True)
            rt_token, src_tokens = src_tokens[0], src_tokens[1:]
            blank_src_smiles = ''.join(src_tokens)
            for i in range(len(src_tokens)):
                src_tokens[i] = add_mapping(src_tokens[i], map_num=i)
            src_smiles = ''.join(src_tokens)
            atom_scores = (batch_atom_scores[0][1:]).cpu().numpy()
            bond_scores = (batch_bond_scores[0][1:, 1:]).cpu().numpy()
            adjacency_matrix = smiles_graph.adjacency_matrix
            full_adjacency_matrix = smiles_graph.full_adjacency_tensor.sum(-1)
            graph_pack = (atom_scores, bond_scores, adjacency_matrix, full_adjacency_matrix)

            cc_trace_with_score_template = get_reaction_centers_from_template(src_smiles, blank_src_smiles,
                                                                              graph_pack, self.rc_candidates)

            if len(cc_trace_with_score_template):
                predicts = [select_diverse_candidate(cc_trace_with_score_template, diverse_k=topk)]
        else:
            predicts = self.infer_reaction_center(atom_rc_scores, bond_rc_scores, graph_packs,
                                                  beta=0.5, percent_aa=0.4, percent_ab=0.55, k=topk,
                                                  num_removal=5, max_count=25)

        decoded_src_token = [self.vocab_src_itos[i] for i in src_token[1:]]

        rc_candidates = []
        for j, (rc, rc_score) in enumerate(predicts[0]):
            src_smiles = self.tag_reaction_center(decoded_src_token, rc)

            if verbose:
                print(rc, rc_score)
                print(src_smiles)

            non_reactive_mask = np.ones(len(decoded_src_token), dtype=bool)
            for idx in rc:
                non_reactive_mask[idx] = 0

            rc_cands = self.build_rc_smarts(cano_prod, non_reactive_mask)
            rc_candidates.append(rc_cands)

        return rc_candidates