import numpy as onp
from scipy.stats import t as student_t
from simtk import unit

from bayes_implicit_solvent.constants import min_r, max_r, min_scale, max_scale
from bayes_implicit_solvent.solvation_free_energy import predict_solvation_free_energy, \
    get_vacuum_samples, create_implicit_sim, beta
from bayes_implicit_solvent.freesolv import db, smiles_list, mol_top_sys_pos_list

use_jax = False
if use_jax:
    import jax.numpy as jnp
    import jax
    from bayes_implicit_solvent.gb_models.jax_gb_models import compute_OBC_energy_vectorized as jax_compute_OBC_energy
import autograd.numpy as np
from bayes_implicit_solvent.gb_models.numpy_gb_models import compute_OBC_energy_vectorized

from bayes_implicit_solvent.solvation_free_energy import kj_mol_to_kT, one_sided_exp

from scipy.spatial.distance import pdist, squareform
from bayes_implicit_solvent.utils import get_charges

# TODO: Maybe resurrect this...
# from functools import lru_cache

class Molecule():
    def __init__(self, smiles, verbose=False, vacuum_samples=None,
                 n_samples=50, thinning=50000, ll='gaussian'):
        """Create an object that supports prediction of solvation free energy given radii and scaling factors

        # TODO: remove ll
        # TODO: add backend attribute

        Parameters
        ----------
        smiles : string
            SMILES representation of molecule

        verbose : boolean
            whether to print updates

        vacuum_samples : list of unit'd snapshots

        n_samples : int
            if vacuum_samples not provided, how many samples to collect?

        thinning : int
            if vacuum samples not provided, how many steps to take between samples?

        ll : string in {'gaussian', 'student-t'}
            is each log-likelihood term Gaussian or heavy-tailed?

        Attributes
        ----------
        mol_index_in_smiles_list : int
            where in the sorted smiles list are we
            TODO: Replace with less janky indexing
        mol_index_in_freesolv : int
            where in freesolv are we
            TODO: Replace with less janky indexing
        mol_name : string
            iupac name (or alternative if IUPAC is unavailable or not parseable by OEChem
        mol : OEMol object
            OEMol object with explicit hydrogens, partial charges, etc. assigned
        top : OpenMM Topology
        sys : OpenMM System
        pos : [n_atoms x 3] array
            initial atomic positions
        atom_names : list of strings
            element symbol + occurrence number, e.g. ['C1', 'F1', 'F2', 'F3', 'Br1']
        n_atoms : int
            number of atoms
        vacuum_sim : OpenMM Simulation
            simulation object at vacuum using Smirnoff parameters, Reference platform, BAOAB Langevin integrator
        vacuum_traj : [n_snapshots x n_atoms x 3] array
            collection of atom positions sampled by vacuum_sim
        experimental_value : float
            experimental value of solvation free energy for this molecule, converted from kcal/mol to unitless (in kT)
        experimental_uncertainty
            uncertainty in experimental_value, converted from kcal.mol to unitless (in kT)
        implicit_sim : OpenMM Simulation
            same as vacuum_sim, but with a GBSAOBCForce added

        """

        assert (ll in {'gaussian', 'student-t'})
        self.smiles = smiles
        self.verbose = verbose
        self.ll = ll

        # find the index of this molecule in our smiles list
        mol_index_in_smiles_list = -1
        for i in range(len(smiles_list)):
            if smiles_list[i] == smiles:
                mol_index_in_smiles_list = i
        self.mol_index_in_smiles_list = mol_index_in_smiles_list
        if self.mol_index_in_smiles_list == -1:
            raise (ValueError(
                "the smiles string queried ({}) doesn't appear to be in FreeSolv's SMILES list...".format(smiles)))

        # find the index of this molecule in Freesolv
        mol_index_in_freesolv = -1
        for i in range(len(db)):
            if db[i][1] == smiles:
                mol_index_in_freesolv = i
        self.mol_index_in_freesolv = mol_index_in_freesolv
        if self.mol_index_in_freesolv == -1:
            raise (ValueError("the smiles string queried ({}) doesn't appear to be in FreeSolv...".format(smiles)))

        self.mol_name = db[mol_index_in_freesolv][2]

        self.mol, self.top, self.sys, self.pos = mol_top_sys_pos_list[self.mol_index_in_smiles_list]
        self.atom_names = [a.name for a in self.top.atoms()]
        self.n_atoms = len(self.pos)

        self._n_samples = n_samples
        self._thinning = thinning
        self.charges = get_charges(self.sys)

        if type(vacuum_samples) == type(None):
            if verbose: print('collecting vacuum samples...')
            self.vacuum_sim, self.vacuum_traj = get_vacuum_samples(self.top, self.sys, self.pos,
                                                                   n_samples=self._n_samples,
                                                                   thinning=self._thinning)
        else:
            self.vacuum_traj = vacuum_samples
            self._n_samples = len(vacuum_samples)

        self.configurations = np.array([snapshot / unit.nanometer for snapshot in self.vacuum_traj])
        self.distance_matrices = np.array([squareform(pdist(snapshot / unit.nanometer)) for snapshot in self.vacuum_traj])

        # both in reduced units
        self.experimental_value = beta * (float(db[mol_index_in_freesolv][3]) * unit.kilocalorie_per_mole)
        self.experimental_uncertainty = beta * (float(db[mol_index_in_freesolv][4]) * unit.kilocalorie_per_mole)

        if verbose:
            print('creating implicit-solvent simulation...')
        self.implicit_sim = create_implicit_sim(self.top, self.sys)

        if verbose:
            print('successfully initialized {}'.format(self.mol_name))

    def predict_solvation_free_energy(self, radii, scaling_factors):
        """Use one-sided EXP to predict the solvation free energy using this set of radii

        Parameters
        ----------
        scaling_factors
        radii : array of floats, either unit'd or assumed to be in nanometers
            radius parameters for each atom
        scaling_factors : array of floats
            scalingFactors for each atom

        Returns
        -------
        mean : float
        uncertainty : float
        """
        assert (len(radii) == self.n_atoms)
        assert (len(scaling_factors) == self.n_atoms)
        return predict_solvation_free_energy(self.implicit_sim, self.vacuum_traj, radii, scaling_factors)

    def predict_solvation_free_energy_autograd(self, radii, scaling_factors):
        W_F = np.array([compute_OBC_energy_vectorized(distance_matrix, radii, scaling_factors, self.charges) for
                    distance_matrix in
                    self.distance_matrices])
        w_F = W_F * kj_mol_to_kT
        return one_sided_exp(w_F)

    #@jit
    def predict_solvation_free_energy_jax(self, radii, scaling_factors):
        @jax.jit
        def compute_component(distance_matrix):
            return jax_compute_OBC_energy(distance_matrix, radii, scaling_factors, self.charges)

        W_F = jax.vmap(compute_component)(self.distance_matrices)

        w_F = W_F * kj_mol_to_kT
        return one_sided_exp(w_F)

    def gaussian_log_likelihood(self, radii, scaling_factors):
        """Un-normalized log-likelihood using Gaussian located at experimentally measured value, with
        scale set by the estimated experimental error.

        This will be sensitive to the stated experimental uncertainty, which we are somewhat skeptical of,
        since for many (most?) entries in FreeSolv this just a default value.

        N(mean_sim | mu = mean_expt,
                     sigma^2 = sigma_expt * max(sigma_expt, sigma_sim)
        """
        simulation_mean, simulation_uncertainty = self.predict_solvation_free_energy(radii, scaling_factors)

        mu = self.experimental_value
        sigma2 = self.experimental_uncertainty * max([self.experimental_uncertainty, simulation_uncertainty])

        return - (simulation_mean - mu) ** 2 / sigma2

    def log_likelihood(self, radii, scaling_factors):
        """To be more robust to inaccurate statement of `experimental_uncertainty`, favor a Student-t likelihood
        over a Gaussian likelihood. This also corresponds to using a nuisance parameter for the experimental uncertainty (with
        an inverse-Gamma prior?) and marginalizing it out.

        TODO: Gelman reference
        """
        simulation_mean, simulation_uncertainty = self.predict_solvation_free_energy(radii, scaling_factors)

        mu = self.experimental_value
        sigma = np.sqrt(self.experimental_uncertainty * max([self.experimental_uncertainty, simulation_uncertainty]))

        return student_t.logpdf(simulation_mean, loc=mu,
                                scale=sigma,  # TODO: Look up how best to put scale information here
                                df=7)  # TODO: Decide what to use for the degrees of freedom parameter

    def bounds_check(self, radii, scale_factors):
        """Check whether parameters are in bounds (radii in (max_r, min_r),
        scale_factors in (min_scale, max_scale)
        """
        # TODO: refactor to have a parameter object that knows how to check its own bounds
        if (np.min(radii) < min_r) or (np.max(radii) > max_r) or \
                (np.min(scale_factors) < min_scale) or (np.max(scale_factors) > max_scale):
            return False
        else:
            return True

    def log_prob_uncached(self, radii, scale_factors):
        """Un-normalized log-probability : log-likelihood, if log-prior > -infty
        """
        if self.bounds_check(radii, scale_factors):
            if self.ll == 'student-t':
                ll = self.log_likelihood(radii, scale_factors)
            elif self.ll == 'gaussian':
                ll = self.gaussian_log_likelihood(radii, scale_factors)
            return ll
        else:
            return - np.inf

    # TODO: Return also a list of all the predictions

    # @lru_cache(maxsize=4)
    # TODO: Maybe use lru_cache again
    def log_prob(self, radii, scale_factors):
        return self.log_prob_uncached(radii, scale_factors)

if __name__ == '__main__':
    from bayes_implicit_solvent.samplers import random_walk_mh

    np.random.seed(0)

    smiles = 'C'
    mol = Molecule(smiles, n_samples=10, thinning=1000)
    n = 2
    radii0 = np.ones(n)
    scaling_factors0 = np.ones(n)
    theta0 = np.hstack([radii0, scaling_factors0])
    assert (len(theta0) == 2 * n)

    def unpack(theta):
        _radii = theta[:n]
        radii = np.zeros(mol.n_atoms)
        radii[0] = _radii[0]
        radii[1:] = _radii[1]

        _scaling_factors = theta[n:]
        scaling_factors = np.zeros(mol.n_atoms)
        scaling_factors[0] = _scaling_factors[0]
        scaling_factors[1:] = _scaling_factors[1]

        return radii, scaling_factors

    def L(theta):
        radii, scaling_factors = unpack(theta)
        return mol.log_prob(radii, scaling_factors)

    traj, log_probs, acceptance_fraction = random_walk_mh(theta0, L,
                                                          n_steps=1000000, stepsize=0.02)
    import os.path

    data_path = 'data/'
    onp.save(os.path.join(data_path, 'radii_samples_{}.npy'.format(smiles)), traj[:, :n])
    onp.save(os.path.join(data_path, 'scale_samples_{}.npy'.format(smiles)), traj[:, n:])
    onp.save(os.path.join(data_path, 'theta_samples_{}.npy'.format(smiles)), traj)
    onp.save(os.path.join(data_path, 'log_probs_{}.npy'.format(smiles)), log_probs)

    print(acceptance_fraction)
    print('atom_names: ', mol.atom_names)

    posterior_predictions = [mol.predict_solvation_free_energy(*unpack(theta))[0] for theta in traj[::100]]

    onp.save(os.path.join(data_path, 'posterior_predictions_{}.npy'.format(smiles)), posterior_predictions)


"""Module for handling Gaussian/CP2K style density files.
"""
import os
from time import time

import numpy as np

from ..utils import fortran_format, python_format, tqdm_wrap

__extensions__ = ['.cube']
__args__ = ['orbitals']

# Unit conversions
bohr_to_ang = .52917721067
ang_to_bohr = 1 / bohr_to_ang

def read(fn, orbitals=0):
    """Read the charge density from a cube file.

    Splits the density into blocks of buffer_size and parses it in chunks.
    Buffer_size must be positive and an integer. If decomposed to orbitals can
    either sum all for total charge density, sum a selection of orbitals or
    return all for user processing. Units are converted to VASP chgcar units.

    args:
        fn: name of the file to open.
        orbitals: how to deal with nval > 1, often associated with molecular
                  orbitals. Passing an iterator returns the sum of the passed
                  orbitals, passing an int > 0 returns just that orbital,
                  passing an int < 0 return the whole charge array in the order
                  density['charge'][nval,nx,ny,nz], the default value, 0,
                  returns a sum of all nvals unless atom number indicator is
                  positive,vin which case it returns just the first nval.
    return:
        density: dict containing 3d-arrays for charge and spin densities.
        lattice: 3x3 array with lattice vectors as rows.
        atoms: atomic positions in Cartesian basis.
        file_info: information about the file type and the write function.
    """
    t0 = time()
    density = dict()
    prefix, filename = os.path.split(fn)
    prefix = os.path.join(prefix, '')
    with open(fn, 'r') as f:
        print(f"  Reading {f.name} as cube format.")
        # first two lines are comments.
        _ = f.readline()
        _ = f.readline()
        line = f.readline().strip().split()
        atom_sum = int(line[0])
        origin = np.array(line[1:4], dtype=np.float64)
        if len(line) > 4:
            nval = int(line[5])
        else:
            nval = 1
        grid = np.zeros(3, dtype=np.int64)
        lattice = np.zeros((3, 3), dtype=np.float64)
        for i in range(3):
            line = f.readline().strip().split()
            grid[i] = line[0]
            lattice[i] = line[1:]
            lattice[i] *= grid[i]
        print(f"  {' x '.join(grid.astype(str))} grid size.")
        # read atomic positions and the blank line after
        atom_types = np.zeros(abs(atom_sum), dtype=np.int64)
        atoms = np.zeros((abs(atom_sum), 3), dtype=np.float64)
        for i in range(abs(atom_sum)):
            line = f.readline().strip().split()
            atom_types[i] = line[0]
            atoms[i] = line[-3:]
        # convert to fractional coordinates and then wrap in cell
        atoms = np.dot(atoms, np.linalg.inv(lattice))
        atoms %= 1
        # convert back
        atoms = np.dot(atoms, lattice)
        if atom_sum < 0:
            line = f.readline().strip().split()
            dset_ids = np.zeros(int(line.pop(0)), dtype=np.int64)
            nval = dset_ids.shape[0]
            count = 0
            while count < nval:
                for m in line:
                    dset_ids[count] = m
                    count += 1
                line = f.readline().strip().split()
        grid_pts = np.prod(grid)
        nx, ny, nz = grid
        record_pts = nz * nval
        grid_lines = record_pts // 6
        grid_mod = record_pts % 6
        # save the current file position and get the line length.
        charge_pos = f.tell()
        line_len = len(f.readline())
        f.seek(charge_pos)
        # set up buffer numbers.
        buffer_size = grid_lines
        buffer_range = [buffer_size]
        charge = np.zeros((nx, ny, nz * nval), dtype=np.float64)
        for x in tqdm_wrap(range(nx), desc="Charge density:"):
            for y in range(ny):
                idx = 0
                for buff in buffer_range:
                    # cube has 6 voxels per line.
                    idx_inc = buff * 6
                    buff_b = buff * line_len
                    line = f.read(buff_b).strip().split()
                    charge[x, y, idx:idx + idx_inc] = line
                    idx += idx_inc
                # get the last non-complete line.
                if grid_mod != 0:
                    line = f.readline().strip().split()
                    charge[x, y, -grid_mod:] = line
        print(f"  File {f.name} closed. ", end='')
    if nval > 1:
        # how are we handling nval > 1
        if hasattr(orbitals, '__iter__'):
            # sum all given orbitals
            charge = charge.reshape(nx, ny, nz, nval)
            charge = np.swapaxes(charge, 0, -1)
            density['charge'] = np.sum([charge[dset_ids.index(int(m))]
                                        for m in orbitals], axis=0)
        elif orbitals < 0:
            # return the entire file
            charge = charge.reshape(nx, ny, nz, nval)
            density['charge'] = np.swapaxes(charge, 0, -1)
        elif orbitals > 0:
            # return specific orbital
            charge = charge.reshape(nx, ny, nz, nval)
            charge = np.swapaxes(charge, 0, -1)
            density['charge'] = charge[dset_ids.index(int(orbitals))].copy()
        elif atom_sum > 0:
            # return just first value (useful for gradient cubes)
            charge = charge.reshape(nx, ny, nz, nval)
            density['charge'] = np.swapaxes(charge, 0, -1)[0].copy()
        else:
            # sum all nvals
            charge = charge.reshape(nx, ny, nz, nval)
            charge = np.swapaxes(charge, 0, -1)
            density['charge'] = np.sum(charge, axis=0)
        del charge
    else:
        density['charge'] = charge
    print(f"Time taken: {time() - t0:0.3f}s", end='\n\n')
    lattice *= bohr_to_ang
    atoms *= bohr_to_ang
    density['charge'] *= ang_to_bohr**3
    file_info = {
        'filename': fn,
        'prefix': prefix,
        'file_type': 'cube',
        'write_function': write,
        'elements': atom_types,
        'voxel_offset': np.array([.5, .5, .5])
    }
    return density, lattice, atoms, file_info

def write(fn, atoms, lattice, density, file_info, prefix=None, suffix='.cube'):
    """Write a cube style charge density

    args:
        fn: filename
        atoms: the atoms for the structure
        lattice: lattice defining cell
        file_info: dictionary containing everything from file_info exported by
                   read function plus optional fortran_format flag
        prefix: string to be placed infront of filename
        suffix: string to be placed at end of filename
    """
    if prefix is not None:
        fn = prefix + fn
    fn += suffix
    if file_info.get('fortran_format', 0) == 2:
        output_format = fortran_format
    elif file_info.get('fortran_format', 0) == 1:
        def output_format(a, p):
            return python_format(a, p, ' ')
    else:
        output_format = python_format
    charge = density['charge']
    # convert to bohr
    atoms *= ang_to_bohr
    charge *= bohr_to_ang**3
    lattice *= ang_to_bohr
    lattice /= charge.shape

    buffer_size = charge.shape[2] // 6
    buffer_rem = charge.shape[2] % 6
    buffer_flag = buffer_rem != 0

    lattice_width = np.max(np.log10(np.abs(lattice[lattice != 0]))) + 9
    lattice_width = max([int(lattice_width), 9]) + 1
    lattice_prec = 17 - lattice_width
    atoms_width = np.max(np.log10(np.abs(atoms[atoms != 0]))) + 9
    atoms_width = max([int(atoms_width), 9]) + 1
    atoms_prec = 17 - atoms_width
    with open(fn, 'w') as f:
        f.write("Cube File writen in pybader\n")
        f.write(file_info['comment'])
        f.write(f"{atoms.shape[0]:>5}{'  0.0000000'*3}\n")
        for i, lat in enumerate(lattice):
            x, y, z = lat
            f.write(f"{charge.shape[i]:>5}")
            f.write(f" {x:> {10}.{lattice_prec}f}")
            f.write(f" {y:> {10}.{lattice_prec}f}")
            f.write(f" {z:> {10}.{lattice_prec}f}\n")
        for i, atom in enumerate(atoms):
            x, y, z = atom
            f.write(f"{file_info['elements'][i]:>5}")
            f.write('  0.0000000')
            f.write(f" {x:> {10}.{atoms_prec}f}")
            f.write(f" {y:> {10}.{atoms_prec}f}")
            f.write(f" {z:> {10}.{atoms_prec}f}\n")
        for i in tqdm_wrap(range(charge.shape[0]), desc=f"{fn}:"):
            for j in range(charge.shape[1]):
                r = charge[i, j][:buffer_size * 6].reshape((buffer_size, 6))
                out = output_format(r, 5)
                if buffer_flag:
                    out += output_format(np.array([charge[i, j]
                                                   [-buffer_rem:]]), 5)
                f.write(out)


import matplotlib as mpl
mpl.use('agg')
import numpy as np, matplotlib.pyplot as plt, os, sys
import tools21cm as t2c, py21cmfast as p21
import matplotlib.gridspec as gridspec
import random, json

from datetime import datetime, date
from glob import glob
from tqdm import tqdm
from mpi4py import MPI

path_out = sys.argv[1]
arr_idx = int(sys.argv[2])
path_out += '/' if path_out[-1]!='/' else ''

path_chache = '/gpfs/scratch/userexternal/mbianco0/21cmFAST-cache/'
p21.config['direc'] = path_chache

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
nprocs = comm.Get_size()

#loop_start, loop_end = 6999, 7000
loop_start, loop_end = np.loadtxt('parameters/todo_r%d.txt' %arr_idx, dtype=int)
perrank = (loop_end-loop_start)//nprocs

def get_dir_size(dir):
    """Returns the "dir" size in bytes."""
    total = 0
    try:
        for entry in os.scandir(dir):
            if entry.is_file():
                total += entry.stat().st_size
            elif entry.is_dir():
                total += get_dir_size(entry.path)
    except NotAdirError:
        return os.path.getsize(dir)
    except PermissionError:
        return 0
    return total

def GenerateSeed():
    # create seed for 21cmFast
    seed = [var for var in datetime.now().strftime('%d%H%M%S')]
    np.random.shuffle(seed)
    return int(''.join(seed))

params = {'HII_DIM':128, 'DIM':384, 'BOX_LEN':256}
c_params = {'OMm':0.27, 'OMb':0.046, 'SIGMA_8':0.82, 'POWER_INDEX':0.96}
my_ext = [0, params['BOX_LEN'], 0, params['BOX_LEN']]
tobs = 1000

if not (os.path.exists('%sastro_params_rank%d.txt' %(path_out+'parameters/', rank))):
    loop_resume = 0

    with open(path_out+'parameters/user_params.txt', 'w') as file:
        file.write(json.dumps(params))

    with open(path_out+'parameters/cosm_params.txt', 'w') as file:
        file.write(json.dumps(c_params))
else:
    loop_resume = int(np.loadtxt('%sastro_params_rank%d.txt' %(path_out+'parameters/', rank))[:,0].max())

if(loop_resume == 0):
    i = int(loop_start+rank*perrank)
elif(loop_resume != 0):
    print(' Rank=%d resumes itration from i=%d.' %(rank, loop_resume))
    i = int(loop_resume + 1)

if(rank == nprocs-1):
    i_end = int(loop_start+(rank+1)*perrank)
else:
    i_end = int(loop_start+(rank+1)*perrank)
    if(i_end != loop_end):
        i_end = loop_end

while i < i_end:
    # astronomical & cosmological parameters
    z = np.random.uniform(7, 9)         # z = [7, 9]
    eff_fact = random.gauss(52.5, 20.)  # eff_fact = [5, 100]
    Rmfp = random.gauss(12.5, 5.)       # Rmfp = [5, 20]
    Tvir = random.gauss(4.65, 0.5)      # Tvir = [log10(1e4), log10(2e5)]

    # Define astronomical parameters
    a_params = {'HII_EFF_FACTOR':eff_fact, 'R_BUBBLE_MAX':Rmfp, 'ION_Tvir_MIN':Tvir}

    # Create 21cmFast cube
    if(i%1 == 0):
        try:
            os.system('rm %s*h5' %path_chache)
        except:
            pass
        comm.Barrier()

        ic = p21.initial_conditions(user_params=params, cosmo_params=c_params, random_seed=GenerateSeed())

    cube = p21.run_coeval(redshift=z, init_box=ic, astro_params=a_params, zprime_step_factor=1.05)

    # Mean neutral fraction
    xn = np.mean(cube.xH_box)

    if(xn > 0.1 and xn <= 0.8):
        print('processor: %d/%d   idx=%d\n z=%.3f, xn=%.3f, eff_fact=%.3f, Rmfp=%.3f, Tvir=%.3f' %(rank, nprocs, i, z, xn, eff_fact, Rmfp, Tvir))

        dT = cube.brightness_temp
        dT1 = t2c.subtract_mean_signal(signal=dT, los_axis=2)

        """
        # calculate uv-coverage
        file_uv, file_Nant = 'uv_coverage_%d/uvmap_z%.3f.npy' %(params['HII_DIM'], z), 'uv_coverage_%d/Nantmap_z%.3f.npy' %(params['HII_DIM'], z)

        if(os.path.exists(file_uv) and os.path.exists(file_Nant)):
            uv = np.load(file_uv)
            Nant = np.load(file_Nant)
        else:
            uv, Nant = t2c.get_uv_daily_observation(params['HII_DIM'], z, filename=None, total_int_time=6.0, int_time=10.0, boxsize=params['BOX_LEN'], declination=-30.0, verbose=True)

            np.save(file_uv, uv)
            np.save(file_Nant, Nant)

        # Noise cube
        np.random.seed(GenerateSeed())
        noise_cube = t2c.noise_cube_coeval(params['HII_DIM'], z, depth_mhz=None, obs_time=tobs, filename=None, boxsize=params['BOX_LEN'], total_int_time=6.0, int_time=10.0, declination=-30.0, uv_map=uv, N_ant=Nant, fft_wrap=False, verbose=False)

        dT2 = dT1 + noise_cube

        # Smooth the data to resolution corresponding to maximum baseline of 2 km
        dT3 = t2c.smooth_coeval(dT2, cube.redshift, box_size_mpc=cube.user_params.HII_DIM, max_baseline=2.0, ratio=1.0, nu_axis=2)

        smt_xn = t2c.smooth_coeval(cube.xH_box, cube.redshift, box_size_mpc=cube.user_params.HII_DIM, max_baseline=2.0, ratio=1.0, nu_axis=2)
        mask_xn = smt_xn>0.5
        """
        # mask are saved such that 1 in neutral region and 0 in ionized region
        t2c.save_cbin(path_out+'data/xH_21cm_i%d.bin' %i, cube.xH_box)
        t2c.save_cbin(path_out+'data/dT1_21cm_i%d.bin' %i, dT1)

        if(i%50):
            ps, ks, n_modes = t2c.power_spectrum_1d(dT, kbins=20, box_dims=cube.user_params.BOX_LEN,return_n_modes=True, binning='log')
            idx=params['HII_DIM']//2

            fig = plt.figure(figsize=(11, 4))
            fig.suptitle('z=%.3f   $x_n$=%.2f   $\zeta$=%.3f   $R_{mfp}$=%.3f   $T_{vir}^{min}$=%.3f' %(z, xn, eff_fact, Rmfp, Tvir), fontsize=15)
            gs = gridspec.GridSpec(nrows=1, ncols=2, width_ratios=[1.5, 1])
            ax0 = fig.add_subplot(gs[0,0])
            ax0.loglog(ks, ps*ks**3/2/np.pi**2)
            ax0.set_xlabel('k (Mpc$^{-1}$)', fontsize=12), ax0.set_ylabel('$\Delta^2_\mathrm{21}$', fontsize=12)
            ax1 = fig.add_subplot(gs[0,1])
            ax1.imshow(dT[:,:,idx], origin='lower', cmap='jet')
            plt.savefig(path_out+'images/test_i%d.png' %i, bbox_inches='tight'), plt.close()

            # Plot outputs comparisons
            fig, axs = plt.subplots(1, 2, figsize=(12,7))
            fig.suptitle('z=%.3f\t\t$x_n$=%.2f\n$\zeta$=%.3f\t\t$R_{mfp}$=%.3f\t\t$T_{vir}^{min}$=%.3f' %(z, xn, eff_fact, Rmfp, Tvir), fontsize=18)
            axs[0].set_title('$x_{HII}$', size=16)
            axs[0].imshow(1-cube.xH_box[:,:,idx], origin='lower', cmap='jet', extent=my_ext)
            axs[0].set_xlabel('[Mpc]'), axs[0].set_ylabel('[Mpc]');

            axs[1].set_title('$\delta T_b$', size=16)
            axs[1].imshow(dT1[:,:,idx], origin='lower', cmap='jet', extent=my_ext)
            axs[1].set_xlabel('[Mpc]'), axs[1].set_ylabel('[Mpc]');
            plt.savefig(path_out+'images/slice_i%d.png' %i, bbox_inches='tight'), plt.close()

        # save parameters values
        with open('%sastro_params_rank%d.txt' %(path_out+'parameters/', rank), 'a') as f:
            if(i == 0 and rank == 0):
                f.write('# HII_EFF_FACTOR: The ionizing efficiency of high-z galaxies\n')
                f.write('# R_BUBBLE_MAX: Mean free path in Mpc of ionizing photons within ionizing regions\n')
                f.write('# ION_Tvir_MIN: Minimum virial Temperature of star-forming haloes in log10 units\n')
                f.write('#i\tz\teff_f\tRmfp\tTvir\tx_n\n')

            f.write('%d\t%.3f\t%.3f\t%.3f\t%.3f\t%.3f\n' %(i, z, eff_fact, Rmfp, Tvir, np.mean(xn)))

        # if output dir is more than 15 GB of size, compress and remove files in data/
        if(get_dir_size(path_out) / 1e9 >= 15):
            if(rank == 0):
                try:
                    strd = np.loadtxt(path_out+'written.txt', dtype=str, delimiter='\n')
                except:
                    strd = np.array([])

                os.system('tar -czvf %s_part%d.tar.gz %s' %(path_out[path_out[:-1].rfind('/')+1:-1], strd.size+1, path_out[path_out[:-1].rfind('/')+1:-1]))
                os.system('rm %sdata/*.bin' %path_out)

                np.savetxt(path_out+'written.txt', np.append(strd, ['%s written %s_part%d.tar.gz' %(datetime.now().strftime('%d/%m/%Y %H:%M:%S'), path_out[path_out[:-1].rfind('/')+1:-1], strd.size+1)]), delimiter='\n', fmt='%s')
            print(' \n Data created exeed 15GB. Compression completed...')
            comm.Barrier()

        # update while loop index
        i += 1
    else:
        continue
comm.Barrier()
print('... rank=%d finished' %rank)


#   Copyright 2020 The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
import contextvars
import functools
import sys
import types
import warnings

from abc import ABCMeta
from functools import singledispatch
from typing import Callable, Iterable, Optional, Sequence, Tuple, Union, cast

import aesara
import numpy as np

from aeppl.logprob import _logcdf, _logprob
from aesara import tensor as at
from aesara.graph.basic import Variable
from aesara.tensor.basic import as_tensor_variable
from aesara.tensor.elemwise import Elemwise
from aesara.tensor.random.op import RandomVariable
from aesara.tensor.var import TensorVariable
from typing_extensions import TypeAlias

from pymc.aesaraf import change_rv_size
from pymc.distributions.shape_utils import (
    Dims,
    Shape,
    Size,
    StrongShape,
    WeakDims,
    convert_dims,
    convert_shape,
    convert_size,
    find_size,
    resize_from_dims,
    resize_from_observed,
)
from pymc.printing import str_for_dist
from pymc.util import UNSET
from pymc.vartypes import string_types

__all__ = [
    "DensityDistRV",
    "DensityDist",
    "Distribution",
    "SymbolicDistribution",
    "Continuous",
    "Discrete",
    "NoDistribution",
]

DIST_PARAMETER_TYPES: TypeAlias = Union[np.ndarray, int, float, TensorVariable]

vectorized_ppc = contextvars.ContextVar(
    "vectorized_ppc", default=None
)  # type: contextvars.ContextVar[Optional[Callable]]

PLATFORM = sys.platform

class _Unpickling:
    pass

class DistributionMeta(ABCMeta):
    """
    DistributionMeta class

    Notes
    -----
    DistributionMeta currently performs many functions, and will likely be refactored soon.
    See issue below for more details
    https://github.com/pymc-devs/pymc/issues/5308
    """

    def __new__(cls, name, bases, clsdict):

        # Forcefully deprecate old v3 `Distribution`s
        if "random" in clsdict:

            def _random(*args, **kwargs):
                warnings.warn(
                    "The old `Distribution.random` interface is deprecated.",
                    FutureWarning,
                    stacklevel=2,
                )
                return clsdict["random"](*args, **kwargs)

            clsdict["random"] = _random

        rv_op = clsdict.setdefault("rv_op", None)
        rv_type = None

        if isinstance(rv_op, RandomVariable):
            rv_type = type(rv_op)

        new_cls = super().__new__(cls, name, bases, clsdict)

        if rv_type is not None:
            # Create dispatch functions

            class_logp = clsdict.get("logp")
            if class_logp:

                @_logprob.register(rv_type)
                def logp(op, values, *dist_params, **kwargs):
                    dist_params = dist_params[3:]
                    (value,) = values
                    return class_logp(value, *dist_params)

            class_logcdf = clsdict.get("logcdf")
            if class_logcdf:

                @_logcdf.register(rv_type)
                def logcdf(op, value, *dist_params, **kwargs):
                    dist_params = dist_params[3:]
                    return class_logcdf(value, *dist_params)

            class_moment = clsdict.get("moment")
            if class_moment:

                @_moment.register(rv_type)
                def moment(op, rv, rng, size, dtype, *dist_params):
                    return class_moment(rv, size, *dist_params)

            # Register the Aesara `RandomVariable` type as a subclass of this
            # `Distribution` type.
            new_cls.register(rv_type)

        return new_cls

def _make_nice_attr_error(oldcode: str, newcode: str):
    def fn(*args, **kwargs):
        raise AttributeError(f"The `{oldcode}` method was removed. Instead use `{newcode}`.`")

    return fn

def _make_rv_and_resize_shape(
    *,
    cls,
    dims: Optional[Dims],
    model,
    observed,
    args,
    **kwargs,
) -> Tuple[Variable, Optional[WeakDims], Optional[Union[np.ndarray, Variable]], StrongShape]:
    """Creates the RV and processes dims or observed to determine a resize shape."""
    # Create the RV without dims information, because that's not something tracked at the Aesara level.
    # If necessary we'll later replicate to a different size implied by already known dims.
    rv_out = cls.dist(*args, **kwargs)
    ndim_actual = rv_out.ndim
    resize_shape = None

    # # `dims` are only available with this API, because `.dist()` can be used
    # # without a modelcontext and dims are not tracked at the Aesara level.
    dims = convert_dims(dims)
    dims_can_resize = kwargs.get("shape", None) is None and kwargs.get("size", None) is None
    if dims is not None:
        if dims_can_resize:
            resize_shape, dims = resize_from_dims(dims, ndim_actual, model)
        elif Ellipsis in dims:
            # Replace ... with None entries to match the actual dimensionality.
            dims = (*dims[:-1], *[None] * ndim_actual)[:ndim_actual]
    elif observed is not None:
        resize_shape, observed = resize_from_observed(observed, ndim_actual)
    return rv_out, dims, observed, resize_shape

class Distribution(metaclass=DistributionMeta):
    """Statistical distribution"""

    rv_class = None
    rv_op: RandomVariable = None

    def __new__(
        cls,
        name: str,
        *args,
        rng=None,
        dims: Optional[Dims] = None,
        initval=None,
        observed=None,
        total_size=None,
        transform=UNSET,
        **kwargs,
    ) -> RandomVariable:
        """Adds a RandomVariable corresponding to a PyMC distribution to the current model.

        Note that all remaining kwargs must be compatible with ``.dist()``

        Parameters
        ----------
        cls : type
            A PyMC distribution.
        name : str
            Name for the new model variable.
        rng : optional
            Random number generator to use with the RandomVariable.
        dims : tuple, optional
            A tuple of dimension names known to the model.
        initval : optional
            Numeric or symbolic untransformed initial value of matching shape,
            or one of the following initial value strategies: "moment", "prior".
            Depending on the sampler's settings, a random jitter may be added to numeric, symbolic
            or moment-based initial values in the transformed space.
        observed : optional
            Observed data to be passed when registering the random variable in the model.
            See ``Model.register_rv``.
        total_size : float, optional
            See ``Model.register_rv``.
        transform : optional
            See ``Model.register_rv``.
        **kwargs
            Keyword arguments that will be forwarded to ``.dist()``.
            Most prominently: ``shape`` and ``size``

        Returns
        -------
        rv : RandomVariable
            The created RV, registered in the Model.
        """

        try:
            from pymc.model import Model

            model = Model.get_context()
        except TypeError:
            raise TypeError(
                "No model on context stack, which is needed to "
                "instantiate distributions. Add variable inside "
                "a 'with model:' block, or use the '.dist' syntax "
                "for a standalone distribution."
            )

        if "testval" in kwargs:
            initval = kwargs.pop("testval")
            warnings.warn(
                "The `testval` argument is deprecated; use `initval`.",
                FutureWarning,
                stacklevel=2,
            )

        if not isinstance(name, string_types):
            raise TypeError(f"Name needs to be a string but got: {name}")

        if rng is None:
            rng = model.next_rng()

        # Create the RV and process dims and observed to determine
        # a shape by which the created RV may need to be resized.
        rv_out, dims, observed, resize_shape = _make_rv_and_resize_shape(
            cls=cls, dims=dims, model=model, observed=observed, args=args, rng=rng, **kwargs
        )

        if resize_shape:
            # A batch size was specified through `dims`, or implied by `observed`.
            rv_out = change_rv_size(rv=rv_out, new_size=resize_shape, expand=True)

        rv_out = model.register_rv(
            rv_out,
            name,
            observed,
            total_size,
            dims=dims,
            transform=transform,
            initval=initval,
        )

        # add in pretty-printing support
        rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
        rv_out._repr_latex_ = types.MethodType(
            functools.partial(str_for_dist, formatting="latex"), rv_out
        )

        rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
        rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)", "pm.logcdf(rv, x)")
        rv_out.random = _make_nice_attr_error("rv.random()", "rv.eval()")
        return rv_out

    @classmethod
    def dist(
        cls,
        dist_params,
        *,
        shape: Optional[Shape] = None,
        size: Optional[Size] = None,
        **kwargs,
    ) -> RandomVariable:
        """Creates a RandomVariable corresponding to the `cls` distribution.

        Parameters
        ----------
        dist_params : array-like
            The inputs to the `RandomVariable` `Op`.
        shape : int, tuple, Variable, optional
            A tuple of sizes for each dimension of the new RV.

            An Ellipsis (...) may be inserted in the last position to short-hand refer to
            all the dimensions that the RV would get if no shape/size/dims were passed at all.
        size : int, tuple, Variable, optional
            For creating the RV like in Aesara/NumPy.

        Returns
        -------
        rv : RandomVariable
            The created RV.
        """
        if "testval" in kwargs:
            kwargs.pop("testval")
            warnings.warn(
                "The `.dist(testval=...)` argument is deprecated and has no effect. "
                "Initial values for sampling/optimization can be specified with `initval` in a modelcontext. "
                "For using Aesara's test value features, you must assign the `.tag.test_value` yourself.",
                FutureWarning,
                stacklevel=2,
            )
        if "initval" in kwargs:
            raise TypeError(
                "Unexpected keyword argument `initval`. "
                "This argument is not available for the `.dist()` API."
            )

        if "dims" in kwargs:
            raise NotImplementedError("The use of a `.dist(dims=...)` API is not supported.")
        if shape is not None and size is not None:
            raise ValueError(
                f"Passing both `shape` ({shape}) and `size` ({size}) is not supported!"
            )

        shape = convert_shape(shape)
        size = convert_size(size)

        create_size, ndim_expected, ndim_batch, ndim_supp = find_size(
            shape=shape, size=size, ndim_supp=cls.rv_op.ndim_supp
        )
        # Create the RV with a `size` right away.
        # This is not necessarily the final result.
        rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)

        # Replicate dimensions may be prepended via a shape with Ellipsis as the last element:
        if shape is not None and Ellipsis in shape:
            replicate_shape = cast(StrongShape, shape[:-1])
            rv_out = change_rv_size(rv=rv_out, new_size=replicate_shape, expand=True)

        rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
        rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)", "pm.logcdf(rv, x)")
        rv_out.random = _make_nice_attr_error("rv.random()", "rv.eval()")
        return rv_out

class SymbolicDistribution:
    """Symbolic statistical distribution

    While traditional PyMC distributions are represented by a single RandomVariable
    graph, Symbolic distributions correspond to a larger graph that contains one or
    more RandomVariables and an arbitrary number of deterministic operations, which
    represent their own kind of distribution.

    The graphs returned by symbolic distributions can be evaluated directly to
    obtain valid draws and can further be parsed by Aeppl to derive the
    corresponding logp at runtime.

    Check pymc.distributions.Censored for an example of a symbolic distribution.

    Symbolic distributions must implement the following classmethods:
    cls.dist
        Performs input validation and converts optional alternative parametrizations
        to a canonical parametrization. It should call `super().dist()`, passing a
        list with the default parameters as the first and only non keyword argument,
        followed by other keyword arguments like size and rngs, and return the result
    cls.num_rngs
        Returns the number of rngs given the same arguments passed by the user when
        calling the distribution
    cls.ndim_supp
        Returns the support of the symbolic distribution, given the default set of
        parameters. This may not always be constant, for instance if the symbolic
        distribution can be defined based on an arbitrary base distribution.
    cls.rv_op
        Returns a TensorVariable that represents the symbolic distribution
        parametrized by a default set of parameters and a size and rngs arguments
    cls.change_size
        Returns an equivalent symbolic distribution with a different size. This is
        analogous to `pymc.aesaraf.change_rv_size` for `RandomVariable`s.
    """

    def __new__(
        cls,
        name: str,
        *args,
        rngs: Optional[Iterable] = None,
        dims: Optional[Dims] = None,
        initval=None,
        observed=None,
        total_size=None,
        transform=UNSET,
        **kwargs,
    ) -> TensorVariable:
        """Adds a TensorVariable corresponding to a PyMC symbolic distribution to the
        current model.

        Parameters
        ----------
        cls : type
            A distribution class that inherits from SymbolicDistribution.
        name : str
            Name for the new model variable.
        rngs : optional
            Random number generator to use for the RandomVariable(s) in the graph.
        dims : tuple, optional
            A tuple of dimension names known to the model.
        initval : optional
            Numeric or symbolic untransformed initial value of matching shape,
            or one of the following initial value strategies: "moment", "prior".
            Depending on the sampler's settings, a random jitter may be added to numeric,
            symbolic or moment-based initial values in the transformed space.
        observed : optional
            Observed data to be passed when registering the random variable in the model.
            See ``Model.register_rv``.
        total_size : float, optional
            See ``Model.register_rv``.
        transform : optional
            See ``Model.register_rv``.
        **kwargs
            Keyword arguments that will be forwarded to ``.dist()``.
            Most prominently: ``shape`` and ``size``

        Returns
        -------
        var : TensorVariable
            The created variable, registered in the Model.
        """

        try:
            from pymc.model import Model

            model = Model.get_context()
        except TypeError:
            raise TypeError(
                "No model on context stack, which is needed to "
                "instantiate distributions. Add variable inside "
                "a 'with model:' block, or use the '.dist' syntax "
                "for a standalone distribution."
            )

        if "testval" in kwargs:
            initval = kwargs.pop("testval")
            warnings.warn(
                "The `testval` argument is deprecated; use `initval`.",
                FutureWarning,
                stacklevel=2,
            )

        if not isinstance(name, string_types):
            raise TypeError(f"Name needs to be a string but got: {name}")

        if rngs is None:
            # Instead of passing individual RNG variables we could pass a RandomStream
            # and let the classes create as many RNGs as they need
            rngs = [model.next_rng() for _ in range(cls.num_rngs(*args, **kwargs))]
        elif not isinstance(rngs, (list, tuple)):
            rngs = [rngs]

        # Create the RV and process dims and observed to determine
        # a shape by which the created RV may need to be resized.
        rv_out, dims, observed, resize_shape = _make_rv_and_resize_shape(
            cls=cls, dims=dims, model=model, observed=observed, args=args, rngs=rngs, **kwargs
        )

        if resize_shape:
            # A batch size was specified through `dims`, or implied by `observed`.
            rv_out = cls.change_size(
                rv=rv_out,
                new_size=resize_shape,
                expand=True,
            )

        rv_out = model.register_rv(
            rv_out,
            name,
            observed,
            total_size,
            dims=dims,
            transform=transform,
            initval=initval,
        )

        # TODO: Refactor this
        # add in pretty-printing support
        rv_out.str_repr = lambda *args, **kwargs: name
        rv_out._repr_latex_ = f"\\text{name}"
        # rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
        # rv_out._repr_latex_ = types.MethodType(
        #     functools.partial(str_for_dist, formatting="latex"), rv_out
        # )

        return rv_out

    @classmethod
    def dist(
        cls,
        dist_params,
        *,
        shape: Optional[Shape] = None,
        size: Optional[Size] = None,
        **kwargs,
    ) -> TensorVariable:
        """Creates a TensorVariable corresponding to the `cls` symbolic distribution.

        Parameters
        ----------
        dist_params : array-like
            The inputs to the `RandomVariable` `Op`.
        shape : int, tuple, Variable, optional
            A tuple of sizes for each dimension of the new RV.
            An Ellipsis (...) may be inserted in the last position to short-hand refer to
            all the dimensions that the RV would get if no shape/size/dims were passed at all.
        size : int, tuple, Variable, optional
            For creating the RV like in Aesara/NumPy.

        Returns
        -------
        var : TensorVariable
        """

        if "testval" in kwargs:
            kwargs.pop("testval")
            warnings.warn(
                "The `.dist(testval=...)` argument is deprecated and has no effect. "
                "Initial values for sampling/optimization can be specified with `initval` in a modelcontext. "
                "For using Aesara's test value features, you must assign the `.tag.test_value` yourself.",
                FutureWarning,
                stacklevel=2,
            )
        if "initval" in kwargs:
            raise TypeError(
                "Unexpected keyword argument `initval`. "
                "This argument is not available for the `.dist()` API."
            )

        if "dims" in kwargs:
            raise NotImplementedError("The use of a `.dist(dims=...)` API is not supported.")
        if shape is not None and size is not None:
            raise ValueError(
                f"Passing both `shape` ({shape}) and `size` ({size}) is not supported!"
            )

        shape = convert_shape(shape)
        size = convert_size(size)

        create_size, ndim_expected, ndim_batch, ndim_supp = find_size(
            shape=shape, size=size, ndim_supp=cls.ndim_supp(*dist_params)
        )
        # Create the RV with a `size` right away.
        # This is not necessarily the final result.
        graph = cls.rv_op(*dist_params, size=create_size, **kwargs)

        # Replicate dimensions may be prepended via a shape with Ellipsis as the last element:
        if shape is not None and Ellipsis in shape:
            replicate_shape = cast(StrongShape, shape[:-1])
            graph = cls.change_size(rv=graph, new_size=replicate_shape, expand=True)

        # TODO: Create new attr error stating that these are not available for DerivedDistribution
        # rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
        # rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)", "pm.logcdf(rv, x)")
        # rv_out.random = _make_nice_attr_error("rv.random()", "rv.eval()")
        return graph

@singledispatch
def _moment(op, rv, *rv_inputs) -> TensorVariable:
    raise NotImplementedError(f"Variable {rv} of type {op} has no moment implementation.")

def moment(rv: TensorVariable) -> TensorVariable:
    """Method for choosing a representative point/value
    that can be used to start optimization or MCMC sampling.

    The only parameter to this function is the RandomVariable
    for which the value is to be derived.
    """
    return _moment(rv.owner.op, rv, *rv.owner.inputs).astype(rv.dtype)

@_moment.register(Elemwise)
def moment_elemwise(op, rv, *dist_params):
    """For Elemwise Ops, dispatch on respective scalar_op"""
    return _moment(op.scalar_op, rv, *dist_params)

class Discrete(Distribution):
    """Base class for discrete distributions"""

    def __new__(cls, name, *args, **kwargs):

        if kwargs.get("transform", None):
            raise ValueError("Transformations for discrete distributions")

        return super().__new__(cls, name, *args, **kwargs)

class Continuous(Distribution):
    """Base class for continuous distributions"""

class NoDistribution(Distribution):
    """Base class for artifical distributions

    RandomVariables that share this type are allowed in logprob graphs
    """

class DensityDistRV(RandomVariable):
    """
    Base class for DensityDistRV

    This should be subclassed when defining custom DensityDist objects.
    """

    name = "DensityDistRV"
    _print_name = ("DensityDist", "\\operatorname{DensityDist}")

    @classmethod
    def rng_fn(cls, rng, *args):
        args = list(args)
        size = args.pop(-1)
        return cls._random_fn(*args, rng=rng, size=size)

class DensityDist(NoDistribution):
    """A distribution that can be used to wrap black-box log density functions.

    Creates a Distribution and registers the supplied log density function to be used
    for inference. It is also possible to supply a `random` method in order to be able
    to sample from the prior or posterior predictive distributions.
    """

    def __new__(
        cls,
        name: str,
        *dist_params,
        logp: Optional[Callable] = None,
        logcdf: Optional[Callable] = None,
        random: Optional[Callable] = None,
        moment: Optional[Callable] = None,
        ndim_supp: int = 0,
        ndims_params: Optional[Sequence[int]] = None,
        dtype: str = "floatX",
        **kwargs,
    ):
        """
        Parameters
        ----------
        name : str
        dist_params : Tuple
            A sequence of the distribution's parameter. These will be converted into
            Aesara tensors internally. These parameters could be other ``RandomVariable``
            instances.
        logp : Optional[Callable]
            A callable that calculates the log density of some given observed ``value``
            conditioned on certain distribution parameter values. It must have the
            following signature: ``logp(value, *dist_params)``, where ``value`` is
            an Aesara tensor that represents the observed value, and ``dist_params``
            are the tensors that hold the values of the distribution parameters.
            This function must return an Aesara tensor. If ``None``, a ``NotImplemented``
            error will be raised when trying to compute the distribution's logp.
        logcdf : Optional[Callable]
            A callable that calculates the log cummulative probability of some given observed
            ``value`` conditioned on certain distribution parameter values. It must have the
            following signature: ``logcdf(value, *dist_params)``, where ``value`` is
            an Aesara tensor that represents the observed value, and ``dist_params``
            are the tensors that hold the values of the distribution parameters.
            This function must return an Aesara tensor. If ``None``, a ``NotImplemented``
            error will be raised when trying to compute the distribution's logcdf.
        random : Optional[Callable]
            A callable that can be used to generate random draws from the distribution.
            It must have the following signature: ``random(*dist_params, rng=None, size=None)``.
            The distribution parameters are passed as positional arguments in the
            same order as they are supplied when the ``DensityDist`` is constructed.
            The keyword arguments are ``rnd``, which will provide the random variable's
            associated :py:class:`~numpy.random.Generator`, and ``size``, that will represent
            the desired size of the random draw. If ``None``, a ``NotImplemented``
            error will be raised when trying to draw random samples from the distribution's
            prior or posterior predictive.
        moment : Optional[Callable]
            A callable that can be used to compute the moments of the distribution.
            It must have the following signature: ``moment(rv, size, *rv_inputs)``.
            The distribution's :class:`~aesara.tensor.random.op.RandomVariable` is passed
            as the first argument ``rv``. ``size`` is the random variable's size implied
            by the ``dims``, ``size`` and parameters supplied to the distribution. Finally,
            ``rv_inputs`` is the sequence of the distribution parameters, in the same order
            as they were supplied when the DensityDist was created. If ``None``, a default
            ``moment`` function will be assigned that will always return 0, or an array
            of zeros.
        ndim_supp : int
            The number of dimensions in the support of the distribution. Defaults to assuming
            a scalar distribution, i.e. ``ndim_supp = 0``.
        ndims_params : Optional[Sequence[int]]
            The list of number of dimensions in the support of each of the distribution's
            parameters. If ``None``, it is assumed that all parameters are scalars, hence
            the number of dimensions of their support will be 0.
        dtype : str
            The dtype of the distribution. All draws and observations passed into the distribution
            will be casted onto this dtype.
        kwargs :
            Extra keyword arguments are passed to the parent's class ``__new__`` method.

        Examples
        --------
            .. code-block:: python

                def logp(value, mu):
                    return -(value - mu)**2

                with pm.Model():
                    mu = pm.Normal('mu',0,1)
                    pm.DensityDist(
                        'density_dist',
                        mu,
                        logp=logp,
                        observed=np.random.randn(100),
                    )
                    idata = pm.sample(100)

            .. code-block:: python

                def logp(value, mu):
                    return -(value - mu)**2

                def random(mu, rng=None, size=None):
                    return rng.normal(loc=mu, scale=1, size=size)

                with pm.Model():
                    mu = pm.Normal('mu', 0 , 1)
                    dens = pm.DensityDist(
                        'density_dist',
                        mu,
                        logp=logp,
                        random=random,
                        observed=np.random.randn(100, 3),
                        size=(100, 3),
                    )
                    prior = pm.sample_prior_predictive(10).prior_predictive['density_dist']
                assert prior.shape == (1, 10, 100, 3)

        """

        if dist_params is None:
            dist_params = []
        elif len(dist_params) > 0 and callable(dist_params[0]):
            raise TypeError(
                "The DensityDist API has changed, you are using the old API "
                "where logp was the first positional argument. In the current API, "
                "the logp is a keyword argument, amongst other changes. Please refer "
                "to the API documentation for more information on how to use the "
                "new DensityDist API."
            )
        dist_params = [as_tensor_variable(param) for param in dist_params]

        # Assume scalar ndims_params
        if ndims_params is None:
            ndims_params = [0] * len(dist_params)

        if logp is None:
            logp = default_not_implemented(name, "logp")

        if logcdf is None:
            logcdf = default_not_implemented(name, "logcdf")

        if moment is None:
            moment = functools.partial(
                default_moment,
                rv_name=name,
                has_fallback=random is not None,
                ndim_supp=ndim_supp,
            )

        if random is None:
            random = default_not_implemented(name, "random")

        rv_op = type(
            f"DensityDist_{name}",
            (DensityDistRV,),
            dict(
                name=f"DensityDist_{name}",
                inplace=False,
                ndim_supp=ndim_supp,
                ndims_params=ndims_params,
                dtype=dtype,
                # Specifc to DensityDist
                _random_fn=random,
            ),
        )()

        # Register custom logp
        rv_type = type(rv_op)

        @_logprob.register(rv_type)
        def density_dist_logp(op, value_var_list, *dist_params, **kwargs):
            _dist_params = dist_params[3:]
            value_var = value_var_list[0]
            return logp(value_var, *_dist_params)

        @_logcdf.register(rv_type)
        def density_dist_logcdf(op, var, rvs_to_values, *dist_params, **kwargs):
            value_var = rvs_to_values.get(var, var)
            return logcdf(value_var, *dist_params, **kwargs)

        @_moment.register(rv_type)
        def density_dist_get_moment(op, rv, rng, size, dtype, *dist_params):
            return moment(rv, size, *dist_params)

        cls.rv_op = rv_op
        return super().__new__(cls, name, *dist_params, **kwargs)

    @classmethod
    def dist(cls, *args, **kwargs):
        output = super().dist(args, **kwargs)
        if cls.rv_op.dtype == "floatX":
            dtype = aesara.config.floatX
        else:
            dtype = cls.rv_op.dtype
        ndim_supp = cls.rv_op.ndim_supp
        return output

def default_not_implemented(rv_name, method_name):
    message = (
        f"Attempted to run {method_name} on the DensityDist '{rv_name}', "
        f"but this method had not been provided when the distribution was "
        f"constructed. Please re-build your model and provide a callable "
        f"to '{rv_name}'s {method_name} keyword argument.\n"
    )

    def func(*args, **kwargs):
        raise NotImplementedError(message)

    return func

def default_moment(rv, size, *rv_inputs, rv_name=None, has_fallback=False, ndim_supp=0):
    if ndim_supp == 0:
        return at.zeros(size, dtype=rv.dtype)
    elif has_fallback:
        return at.zeros_like(rv)
    else:
        raise TypeError(
            "Cannot safely infer the size of a multivariate random variable's moment. "
            f"Please provide a moment function when instantiating the {rv_name} "
            "random variable."
        )


# -*-coding:utf-8 -*-

'''
@File       : beam_search.py
@Author     : HW Shen
@Date       : 2020/8/12
@Desc       :
'''

import numpy as np
import tensorflow as tf

BATCH_SIZE = 1  # 每个batch数据量
BEAM_SIZE = 2  # beam_search的尺度
EMB_DIM = 10  # 词向量维度
ENCODER_UNITS = 20  # encoder层的单位数
DECODER_UNITS = 20  # decoder层的单位数
ATTENTION_UNITS = 20  # attention层的单位数
MAX_LEN = 10  # Decoding输出的最大长度
MIN_LEN = 1   # Decoding输出的最小长度
START_TOKEN = 'START'  # Decoder过程的开始标识符
END_TOKEN = 'END'  # Deocder过程的结束标识符
UNK_TOKEN = 'UNK'  # OOV

word2id = {'START': 0, 'END': 1, 'PAD': 2, '我': 3, '你': 4, '洗澡': 5, '吃饭': 6, 'UNK': 7}
id2word = ['START', 'END', 'PAD', '我', '你', '洗澡', '吃饭', 'UNK']

class Encoder(tf.keras.layers.Layer):
    """编码器 : 双层GRU """

    def __init__(self, vocab_size, embedding_dim, encoder_units, batch_size):

        super(Encoder, self).__init__()
        self.batch_size = batch_size
        self.encoder_units = encoder_units // 2  # 因为是双向的，所以每一层均分encoder_units
        self.embedding = tf.keras.layers.Embedding(input_dim=vocab_size, output_dim=embedding_dim)
        self.gru = tf.keras.layers.GRU(units=self.encoder_units,
                                       return_sequences=True,
                                       return_state=True,
                                       recurrent_initializer='glorot_uniform')
        self.bi_gru = tf.keras.layers.Bidirectional(layer=self.gru, merge_mode='concat')

    def call(self, x, hidden):
        x = self.embedding(x)
        hidden = tf.split(value=hidden, num_or_size_splits=2, axis=1)
        output, forward_state, backward_state = self.bi_gru(x, initial_state=hidden)
        state = tf.concat([forward_state, backward_state], axis=1)

        return output, state

    def initialize_hidden_state(self):
        # encoder隐层初始化参数
        return tf.zeros((self.batch_size, self.encoder_units * 2))  # shape:(batch_size, encoder_units)

class Attention(tf.keras.layers.Layer):
    """attention 类"""

    def __init__(self, units):
        super(Attention, self).__init__()
        self.W_s = tf.keras.layers.Dense(units)
        self.W_h = tf.keras.layers.Dense(units)
        self.V = tf.keras.layers.Dense(1)

    def call(self, decoder_hidden, encoder_output):

        hidden_with_time_axis = tf.expand_dims(decoder_hidden, 1)  # 给dec_hidden增加一个时间维度
        score = self.V(tf.nn.tanh(self.W_s(encoder_output) + self.W_h(hidden_with_time_axis)))  # 注意力得分
        attn_dist = tf.nn.softmax(score, axis=1)  # 注意力得分归一化
        context_vector = attn_dist * encoder_output  # 输入词向量的上下文向量
        context_vector = tf.reduce_sum(context_vector, axis=1)  # 上下文向量的求和，降维

        return context_vector, tf.squeeze(attn_dist, -1)

class Decoder(tf.keras.layers.Layer):
    """ 解码层 ： 单层 GRU + Dense"""

    def __init__(self, vocab_size, embedding_dim, decoder_units, batch_size):
        super(Decoder, self).__init__()
        self.batch_size = batch_size
        self.decoder_units = decoder_units  # Decoder是单层GRU
        self.embedding = tf.keras.layers.Embedding(input_dim=vocab_size, output_dim=embedding_dim)
        self.gru = tf.keras.layers.GRU(units=self.decoder_units,
                                       return_sequences=True,
                                       return_state=True,
                                       recurrent_initializer='glorot_uniform')
        self.fc = tf.keras.layers.Dense(units=vocab_size, activation=tf.keras.activations.softmax)

    def call(self, x, hidden, enc_output, context_vector):
        x = self.embedding(x)
        x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)
        output, state = self.gru(x)
        output = tf.reshape(output, (-1, output.shape[2]))
        out = self.fc(output)
        return x, out, state

class Seq2SeqModel(object):
    """seq2seq 类"""

    def __init__(self, encoder, decoder, attention):

        self.encoder = encoder
        self.decoder = decoder
        self.attention = attention

    def call_encoder(self, encoder_input):

        encoder_hidden = self.encoder.initialize_hidden_state()  # 初始化隐层参数矩阵
        encoder_output, encoder_hidden = self.encoder(encoder_input, encoder_hidden)  # 编码器的 output 和 hidden_state

        return encoder_output, encoder_hidden

    def decode_one_step(self, latest_tokens, encoder_states, decoder_in_states):
        """
        Decoder层的每一步解码
        Args:
            latest_tokens: 当前时间步, 解码器输入的 token
            encoder_states: 编码器输出的 output
            decoder_in_states: 解码器上一时间步传来的隐层向量
        Returns:
            top_k_ids: [beam, beam * 2], 解码结果中 top 2 * beam 的 token id
            top_k_log_probs: [beam, beam * 2], 解码过程中 top 2 * beam 的 log 概率得分
            new_states: list, 每个 item 是每个 beam 中的隐层向量
        """
        # latest_tokens = np.transpose(np.array([latest_tokens]))
        # decoder_in_states = np.concatenate(arrays=decoder_in_states, axis=0)
        latest_tokens = tf.transpose([latest_tokens])
        decoder_in_states = tf.concat(values=decoder_in_states, axis=0)
        context_vector, attn_dists = self.attention(decoder_in_states, encoder_states)
        _, prediction, decoder_hidden = self.decoder(latest_tokens, decoder_in_states, encoder_states, context_vector)
        # 把上一步预测的结果prediction作为topK的选依据
        top_k_log_probs, top_k_ids = tf.nn.top_k(input=prediction, k=BEAM_SIZE * 2, sorted=True)
        top_k_log_probs = tf.math.log(x=top_k_log_probs)
        new_states = [np.expand_dims(decoder_hidden[i, :], axis=0) for i in range(BEAM_SIZE)]

        return top_k_ids.numpy(), top_k_log_probs, new_states, attn_dists

class BeamHypotheses(object):
    """ BeamSearch的假设 """

    def __init__(self, token_list, log_prob_list, state, attn_dist_list):

        self.token_list = token_list  # list of all the tokens from time 0 to the current time step t
        self.log_prob_list = log_prob_list  # list of the log probabilities of the tokens
        self.state = state  # decoder state after the last token decoding
        self.attn_dist_list = attn_dist_list  # attention dists of all the tokens

    def extend(self, token, log_prob, state, attn_dist):
        """
        Method to extend the current hypothesis by adding the next decoded token
        Args:
            token: the next decoded token
            log_prob: the log prob of the next decoded token
            state: next hidden state
            attn_dist: # the attention dist of the next decoded toke
        Returns: next hypothesis
        """
        return BeamHypotheses(
            token_list=self.token_list + [token],
            log_prob_list=self.log_prob_list + [log_prob],
            state=state,
            attn_dist_list=self.attn_dist_list + [attn_dist]
        )

    @property
    def latest_token(self):
        return self.token_list[-1]

    @property
    def log_prob(self):
        return sum(self.log_prob_list)

    @property
    def avg_log_prob(self):
        return self.log_prob / len(self.token_list)

class Generation(object):
    """ 按照beam_search的方式 生成Decode的相应结果 """

    def __init__(self):
        pass

    @staticmethod
    def sort_hyp(hyp_list):
        # 对所有假设从大到小排序（按照平均概率值）
        return sorted(hyp_list, key=lambda h: h.avg_log_prob, reverse=True)

    def generate_beam_search(self, model, model_input):

        # encoder states bi-gru 的 output; decoder in state 是拼接的隐层向量
        encoder_states, decoder_in_state = model.call_encoder(model_input)
        hyp_list = [
            BeamHypotheses(token_list=[word2id[START_TOKEN]],
                           log_prob_list=[0.0],
                           state=decoder_in_state,
                           attn_dist_list=[],
                           ) for _ in range(BEAM_SIZE)
        ]
        result = []
        step = 0

        while step < MAX_LEN and len(result) < BEAM_SIZE:
            # 容器里每个 hpy 的最后一个 token 的列表
            latest_tokens = [hyp.latest_token for hyp in hyp_list]
            # 处理异常, 搞成 unk 的 id
            latest_tokens = [token if token in range(len(id2word)) else word2id[UNK_TOKEN] for token in latest_tokens]
            # 容器里的 hpy 在解码器得到 state 的列表
            states = [hyp.state for hyp in hyp_list]

            top_k_ids, top_k_log_probs, new_states, attn_dists = model.decode_one_step(latest_tokens=latest_tokens,
                                                                                encoder_states=encoder_states,
                                                                                decoder_in_states=states)
            # extend 操作后, 都收集在 all_hyp_list 中
            all_hyp_list = []
            # 初始状态开始 beam search 时, 实际只有一种状态; 但在后面的 step 时, 每次都有 beam 个状态
            num_ori_hyp = 1 if step == 0 else len(hyp_list)

            for i in range(num_ori_hyp):
                hyp = hyp_list[i]
                new_state = new_states[i]
                attn_dist = attn_dists[i]

                for j in range(BEAM_SIZE * 2):
                    new_hpy = hyp.extend(token=top_k_ids[i, j],
                                         log_prob=top_k_log_probs[i, j],
                                         state=new_state,
                                         attn_dist=attn_dist)
                    all_hyp_list.append(new_hpy)

            # 对 hpy 进行排序, 只有满足了 END TOKEN 和 MIN LEN 才会加入result
            hyp_list = []
            for hyp in self.sort_hyp(all_hyp_list):
                # 出现终止符号
                if hyp.latest_token == word2id[END_TOKEN]:
                    if step >= MIN_LEN:
                        result.append(hyp)
                else:
                    hyp_list.append(hyp)
                if len(hyp_list) == BEAM_SIZE or len(result) == BEAM_SIZE:
                    break

            step += 1

        # 如果循环都结束了, 但都仍然没有 END 出现, result 就是一个空 hyp_list, 此时要把此时的 hyp list 进行排序取最大的输出
        if len(result) == 0:
            result = hyp_list
        hyp_sorted = self.sort_hyp(result)

        return ' '.join(id2word[index] for index in hyp_sorted[0].token_list)

if __name__ == '__main__':
    model_obj = Seq2SeqModel(encoder=Encoder(vocab_size=len(word2id),
                                             embedding_dim=EMB_DIM,
                                             encoder_units=ENCODER_UNITS,
                                             batch_size=BATCH_SIZE),
                             decoder=Decoder(vocab_size=len(word2id),
                                             embedding_dim=EMB_DIM,
                                             decoder_units=DECODER_UNITS,
                                             batch_size=BATCH_SIZE),
                             attention=Attention(units=ATTENTION_UNITS))
    generator_obj = Generation()
    res = generator_obj.generate_beam_search(model_obj, np.array([[0, 3, 5, 5, 5]]))
    print(res)


#!/usr/bin/env python
#
# Copyright 2019 DFKI GmbH.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the
# following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
# USE OR OTHER DEALINGS IN THE SOFTWARE.
import numpy as np
from .fabrik_chain import FABRIKChain, to_local_cos, quaternion_from_vector_to_vector, FABRIKBone, ROOT_OFFSET

def get_child_offset(skeleton, node, child_node):
    actual_offset = np.array([0, 0, 0], np.float)
    while node is not None and skeleton.nodes[node].children[0].node_name != child_node:
        local_offset = np.array(skeleton.nodes[node].children[0].offset)
        local_offset_len = np.linalg.norm(local_offset)
        if local_offset_len > 0:
            local_offset /= local_offset_len
            actual_offset = actual_offset + local_offset
            node = skeleton.nodes[node].children[0].node_name
            if len(skeleton.nodes[node].children) < 1:
                node = None
        else:
            node = None
    return actual_offset

def get_global_joint_parameters_from_positions(skeleton, positions, next_nodes):
    n_parameters = len(skeleton.animated_joints)*4+3
    frame = np.zeros(n_parameters)
    frame[:3] = positions[skeleton.root]
    o = 3
    print(skeleton.animated_joints)
    print(next_nodes.keys())
    print(positions.keys())
    for node in skeleton.animated_joints:
        next_node = next_nodes[node]
        if next_node in positions:
            next_offset = np.array(skeleton.nodes[next_node].offset)
            next_offset /= np.linalg.norm(next_offset)

            # 1. sum over offsets of static nodes
            local_offset = get_child_offset(skeleton, node, next_node)
            #print(node, local_offset)
            actual_offset = next_offset + local_offset
            actual_offset /= np.linalg.norm(actual_offset)  # actual_offset = [0.5, 0.5,0]

            dir_to_child = positions[next_node] - positions[node]
            dir_to_child /= np.linalg.norm(dir_to_child)

            # 2. get global rotation
            global_q = quaternion_from_vector_to_vector(actual_offset, dir_to_child)
            frame[o:o + 4] = global_q
        else:
            #print("ignore", node, next_node)
            frame[o:o + 4] = [1,0,0,0]
        o += 4
    return frame

def global_to_local_frame(skeleton, global_frame):
    n_parameters = len(skeleton.animated_joints)*4+3
    frame = np.zeros(n_parameters)
    o = 3
    for node in skeleton.animated_joints:
        frame[o:o + 4] = to_local_cos(skeleton, node, frame, global_frame[o:o + 4])

        o += 4

    return frame

class FABRIKNode(object):
    def __init__(self, skeleton, root, parent_chain=None, tolerance=0.01, max_iter=100, root_offset=ROOT_OFFSET):
        self.root_pos = np.array([0,0,0], dtype=np.float)
        self.skeleton = skeleton
        self.n_parameters = len(self.skeleton.animated_joints)*4+3
        self.tolerance = tolerance
        self.max_iter = max_iter
        self.parent_chain = parent_chain
        self.root = root
        self.root_offset = root_offset
        self.child_nodes = list()
        # self.target = None
        # self.position = None
        self.is_leaf = True
        self.construct_from_skeleton(skeleton, root, tolerance, max_iter)

    def construct_from_skeleton(self, skeleton, root, tolerance, max_iter):
        """ there need to be two end effectors"""
        children = [c.node_name for c in skeleton.nodes[root].children]
        for c in children:
            self.is_leaf = False
            current_node = c
            node_order = [self.root]
            while current_node is not None:
                n_children = len(skeleton.nodes[current_node].children)
                if n_children == 1:  # append to chain
                    child_node = skeleton.nodes[current_node].children[0].node_name
                    # only add list to joints
                    if not skeleton.nodes[current_node].fixed:
                        node_order.append(current_node)# skip fixed nodes
                    current_node = child_node
                else:  # stop chain # split up by adding child nodes
                    if n_children > 0:
                        node_order.append(current_node)
                    bones = dict()
                    for idx, node in enumerate(node_order):
                        child_node = None
                        if idx+1 < len(node_order):
                            child_node = node_order[idx + 1]
                        bones[node] = FABRIKBone(node, child_node)
                        if idx == 0 and self.parent_chain is None :
                            bones[node].is_root = True
                        else:
                            bones[node].is_root = False
                    parent_chain = FABRIKChain(skeleton, bones, node_order)
                    print("construct node at",self.root , current_node, node_order)
                    node = FABRIKNode(skeleton, current_node, parent_chain, tolerance, max_iter)
                    self.child_nodes.append(node)
                    current_node = None

    def backward_stage(self):
        for c in self.child_nodes:
            c.backward_stage()  # calculate backward update of children of child to start at the end effectors

        if not self.is_leaf:
            positions = [c.get_chain_root_position() for c in self.child_nodes]
            t = np.mean(positions, axis=0)
            #print("centroid",t)
            #print("no leaf")
        else:
            t = self.parent_chain.target
            #print("leaf")

        if self.parent_chain is not None:
            self.parent_chain.target = t
            if self.target_is_reachable():
               self.parent_chain.backward()
            else:
                print("unreachable")
                # if unreachable orient joints to target
                self.parent_chain.orient_to_target()

    def target_is_reachable(self):
        return self.parent_chain.target_is_reachable()

    def forward_stage(self):
        # calculate forward update of parent chain
        if self.parent_chain is not None:
            #if self.parent_chain.target_is_reachable():
            self.parent_chain.forward()
            #else:
            #    self.parent_chain.orient_to_target()

        # calculate forward update of children of child
        for c in self.child_nodes:
            if self.parent_chain is not None:
                c.parent_chain.root_pos = self.parent_chain.get_end_effector_position()
            else:
                c.parent_chain.root_pos = self.root_pos
            c.forward_stage()

    def get_chain_root_position(self):
        root_node = self.parent_chain.node_order[0]
        return self.parent_chain.bones[root_node].position

    def set_positions_from_frame(self, frame, parent_length):
        if self.parent_chain is not None:
            self.parent_chain.set_positions_from_frame(frame, parent_length)
            parent_length += self.parent_chain.chain_length
        for c in self.child_nodes:
            c.set_positions_from_frame(frame, parent_length)

    def get_error(self):
        error = 0
        for c in self.child_nodes:
            error += c.get_error()
        if self.is_leaf:
            if self.parent_chain.target_is_reachable():
                error += self.parent_chain.get_error()
            else:
                error += 0
        return error

    def solve(self):
        iter = 0
        distance = self.get_error()
        while distance > self.tolerance and iter < self.max_iter:
            self.backward_stage()
            self.forward_stage()
            distance = self.get_error()
            iter += 1
        n_out_of_reach = self.targets_out_of_reach()
        if n_out_of_reach > 0:
            self.orient_to_targets()
        print("solved", distance, n_out_of_reach)
        self.print_end_effectors()

    def print_end_effectors(self):
        n_targets = 0
        for c in self.child_nodes:
            if c.is_leaf:
                print(c.root, c.parent_chain.target, c.parent_chain.get_end_effector_position())
            else:
                c.print_end_effectors()
        return n_targets

    def targets_out_of_reach(self):
        n_targets = 0
        for c in self.child_nodes:
            if c.is_leaf and not c.parent_chain.target_is_reachable():
                n_targets+=1
            else:
                n_targets += c.targets_out_of_reach()
        return n_targets

    def orient_to_targets(self):
        for c in self.child_nodes:
            if c.is_leaf and not c.parent_chain.target_is_reachable():
                c.parent_chain.orient_to_target()
            else:
                c.orient_to_targets()

    def get_joint_parameters(self, frame):
        #print("get joint parameters")
        for c in self.child_nodes:
            print("from", c.parent_chain.node_order[0], "to", c.parent_chain.node_order[-1])
            global_frame = c.parent_chain.get_joint_parameters_global()
            src = 3
            if self.parent_chain is None:
                animated_joints = c.parent_chain.node_order
            else:
                animated_joints = c.parent_chain.node_order[1:]
            for j in animated_joints: #ignore root rotation
                dest = self.skeleton.animated_joints.index(j)*4 +3
                frame[dest:dest+4] = to_local_cos(self.skeleton, j, frame, global_frame[src:src+4])
                src += 4
            c.get_joint_parameters(frame)
        return frame

    def set_targets(self, targets):
        #print("set targets",targets)
        if self.is_leaf:
            if self.root in targets:
                self.target = targets[self.root]
                self.parent_chain.target = self.target
                #print("set target",self.root)
            else:
                self.target = self.parent_chain.get_end_effector_position() # keep the initial position as target for unconstrained leafs
                self.parent_chain.target = self.target
                print("set end effector position")
        else:
            for c in self.child_nodes:
                c.set_targets(targets)

    def run(self, orig_frame, targets):
        self.set_positions_from_frame(orig_frame, 0)
        self.set_targets(targets)
        self.solve()
        if False:
            frame = np.zeros(self.n_parameters)
            self.get_joint_parameters(frame)
        else:
            joint_pos = dict()
            self.get_global_positions(joint_pos)
            next_nodes = dict()
            self.get_next_nodes(next_nodes)
            frame = get_global_joint_parameters_from_positions(self.skeleton, joint_pos, next_nodes)
            frame = global_to_local_frame(self.skeleton, frame)
            #print("done")
            #self.print_frame_parameters(frame)
        return frame

    def get_next_nodes(self, next_nodes):
        if self.parent_chain is not None:
            self.parent_chain.get_next_nodes(next_nodes)
        for c in self.child_nodes:
            c.get_next_nodes(next_nodes)

    def print_frame_parameters(self, frame):
        idx = 3
        for n in self.skeleton.nodes:
            if len(self.skeleton.nodes[n].children) > 0:
                print(n, frame[idx: idx+4])
                idx +=4
        return

    def draw(self, m,v,p, l):
        if self.parent_chain is not None:
            self.parent_chain.draw(m,v,p,l)
        for c in self.child_nodes:
            c.draw(m, v, p, l)

    def get_global_positions(self, positions):
        if self.parent_chain is not None:
            positions.update(self.parent_chain.get_global_positions())
        for c in self.child_nodes:
            c.get_global_positions(positions)