import random
from pathlib import Path
import math
import logging
import os

from collections import defaultdict
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

import gym
from gym import spaces
from gym.spaces.box import Box
from gym.spaces.discrete import Discrete

import numpy as np

from deep_rl import Task, Config, Logger, BaseNormalizer
from deep_rl.component.envs import DummyVecEnv, SubprocVecEnv, OriginalReturnWrapper
from deep_rl.network import GaussianActorCriticNet, DeterministicActorCriticNet, NoisyLinear, layer_init

from py_diff_pd.core.py_diff_pd_core import HexMesh3d, HexDeformable, StdRealVector, StdIntVector, QuadMesh2d
from py_diff_pd.common.common import create_folder, ndarray, print_info
from py_diff_pd.common.hex_mesh import generate_hex_mesh, get_boundary_face
from py_diff_pd.common.display import export_gif, Arrow3D

from baselines.common.running_mean_std import RunningMeanStd

class MeanStdNormalizer(BaseNormalizer):
    def __init__(self, read_only=False, clip=10.0, epsilon=1e-8):
        BaseNormalizer.__init__(self, read_only)
        self.read_only = read_only
        self.rms = None
        self.clip = clip
        self.epsilon = epsilon

    def __call__(self, x):
        x = np.asarray(x)
        if self.rms and self.read_only:
            return np.clip(x, -self.clip, self.clip)
        if self.rms is None:
            self.rms = RunningMeanStd(shape=(1,) + x.shape[1:])
        if not self.read_only:
            self.rms.update(x)
        return np.clip((x - self.rms.mean) / np.sqrt(self.rms.var + self.epsilon),
                       -self.clip, self.clip)

    def state_dict(self):
        return {'mean': self.rms.mean,
                'var': self.rms.var}

    def load_state_dict(self, saved):
        self.rms.mean = saved['mean']
        self.rms.var = saved['var']

def tensor(x):
    if isinstance(x, torch.Tensor):
        return x
    x = np.asarray(x, dtype=np.float64)
    x = torch.from_numpy(x).to(Config.DEVICE)
    return x

class LayerNormFCBody(nn.Module):
    def __init__(self, state_dim, hidden_units=(64, 64), gate=F.relu, noisy_linear=False):
        super(LayerNormFCBody, self).__init__()
        dims = (state_dim,) + hidden_units
        if noisy_linear:
            self.layers = nn.ModuleList(
                [NoisyLinear(dim_in, dim_out) for dim_in, dim_out in zip(dims[:-1], dims[1:])])
        else:
            self.layers = nn.ModuleList(
                [nn.Sequential(*[
                    nn.Linear(dim_in, dim_out, bias=False),
                    nn.LayerNorm(dim_out, elementwise_affine=True)
                    ]) for dim_in, dim_out in zip(dims[:-1], dims[1:])])

        self.gate = gate
        self.feature_dim = dims[-1]
        self.noisy_linear = noisy_linear

    def reset_noise(self):
        if self.noisy_linear:
            for layer in self.layers:
                layer.reset_noise()

    def forward(self, x):
        for layer in self.layers:
            x = self.gate(layer(x))
        return x

class MyDeterministicActorCriticNet(DeterministicActorCriticNet):
    def feature(self, obs):
        obs = tensor(obs)
        return self.phi_body(obs)

    def save(self, path):
        torch.save({
            'checkpoints': self.network.state_dict(),
            'normalizer': self.config.state_normalizer.state_dict(),
        }, path)

    def load(self, path):
        state_dict = torch.load(path, map_location='cpu')
        self.network.load_state_dict(state_dict['checkpoints'])
        self.config.state_normalizer.load_state_dict(state_dict['normalizer'])

class MyGaussianActorCriticNet(GaussianActorCriticNet):
    def forward(self, obs, action=None):
        obs = tensor(obs)
        phi = self.phi_body(obs)
        phi_a = self.actor_body(phi)
        phi_v = self.critic_body(phi)
        mean = torch.tanh(self.fc_action(phi_a))
        v = self.fc_critic(phi_v)
        dist = torch.distributions.Normal(mean, F.softplus(self.std))
        if action is None:
            action = dist.sample()
        log_prob = dist.log_prob(action).sum(-1).unsqueeze(-1)
        entropy = dist.entropy().sum(-1).unsqueeze(-1)
        return {'action': action,
                'log_pi_a': log_prob,
                'entropy': entropy,
                'mean': mean,
                'v': v}

    def save(self, path):
        torch.save({
            'checkpoints': self.network.state_dict(),
            'normalizer': self.config.state_normalizer.state_dict(),
        }, path)

    def load(self, path):
        state_dict = torch.load(path, map_location='cpu')
        self.network.load_state_dict(state_dict['checkpoints'])
        self.config.state_normalizer.load_state_dict(state_dict['normalizer'])

logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s: %(message)s')

def get_logger(path, log_level=0):
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    return Logger(logger, str(path), log_level)

class DiffPDTask(Task):
    def __init__(
            self,
            env_fn,
            sim_class,
            seed,
            num_envs=1,
            single_process=True,
            episode_life=True,
        ):

        envs = [make_env(env_fn, sim_class, seed, i, episode_life) for i in range(num_envs)]
        if num_envs == 1 or single_process:
            Wrapper = DummyVecEnv
        else:
            Wrapper = SubprocVecEnv
        self.env = Wrapper(envs)
        self.name = 'diffpd'
        self.observation_space = self.env.observation_space
        self.state_dim = int(np.prod(self.env.observation_space.shape))

        self.action_space = self.env.action_space
        if isinstance(self.action_space, Discrete):
            self.action_dim = self.action_space.n
        elif isinstance(self.action_space, Box):
            self.action_dim = self.action_space.shape[0]
        else:
            assert 'unknown action space'

    def reset(self):
        return self.env.reset()

    def step(self, actions):
        if isinstance(self.action_space, Box):
            actions = np.clip(actions, self.action_space.low, self.action_space.high)
        return self.env.step(actions)

def make_env(env_fn, *args, **kwargs):
    def _thunk():
        env = env_fn(*args, **kwargs)
        env = OriginalReturnWrapper(env)

        return env
    return _thunk

def make_water_snake_3d(sim_class, seed, rank, *args, **kwargs):

    os.system("taskset -p 0xff %d" % os.getpid())

    seed = seed + rank
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.set_default_dtype(torch.float64)

    folder = Path('water_snake').resolve()
    folder.mkdir(parents=True, exist_ok=True)

    # Mesh parameters.
    cell_nums = [20, 2, 2]
    node_nums = [c + 1 for c in cell_nums]
    dx = 0.1
    origin = np.zeros((3,))
    bin_file_name = str(folder / 'water_snake.bin')
    voxels = np.ones(cell_nums)

    voxel_indices, vertex_indices = generate_hex_mesh(voxels, dx, origin, bin_file_name, write=False)
    mesh = HexMesh3d()
    mesh.Initialize(bin_file_name)

    # FEM parameters.
    youngs_modulus = 1e6
    poissons_ratio = 0.45
    density = 1e3
    method = 'pd_eigen'
    opt = {
        'max_pd_iter': 1000, 'max_ls_iter': 10, 'abs_tol': 1e-4, 'rel_tol': 1e-3, 'verbose': 0,
        'thread_ct': 1, 'use_bfgs': 1, 'bfgs_history_size': 10
    }

    deformable = HexDeformable()
    deformable.Initialize(bin_file_name, density, 'none', youngs_modulus, poissons_ratio)
    # Elasticity.
    deformable.AddPdEnergy('corotated', [youngs_modulus / (1 + poissons_ratio),], [])
    # Hydrodynamics parameters.
    rho = 1e3
    v_water = [0, 0, 0]   # Velocity of the water.
    # # Cd_points = (angle, coeff) pairs where angle is normalized to [0, 1].
    Cd_points = ndarray([[0.0, 0.05], [0.4, 0.05], [0.7, 1.85], [1.0, 2.05]])
    # # Ct_points = (angle, coeff) pairs where angle is normalized to [-1, 1].
    Ct_points = ndarray([[-1, -0.8], [-0.3, -0.5], [0.3, 0.1], [1, 2.5]])
    # The current Cd and Ct are similar to Figure 2 in SoftCon.
    # surface_faces is a list of (v0, v1) where v0 and v1 are the vertex indices of the two endpoints of a boundary edge.
    # The order of (v0, v1) is determined so that following all v0 -> v1 forms a ccw contour of the deformable body.
    surface_faces = get_boundary_face(mesh)
    deformable.AddStateForce(
        'hydrodynamics', np.concatenate(
            [[rho,], v_water, Cd_points.ravel(), Ct_points.ravel(), ndarray(surface_faces).ravel()]))

    # Add actuation.
    # ******************** <- muscle
    # |                  | <- body
    # |                  | <- body
    # ******************** <- muscle

    all_muscles = []
    shared_muscles = []
    for i in [0, cell_nums[2] - 1]:
        muscle_pair = []
        for j in [0, cell_nums[1] - 1]:
            indices = voxel_indices[:, j, i].tolist()
            deformable.AddActuation(1e5, [1.0, 0.0, 0.0], indices)
            muscle_pair.append(indices)
        shared_muscles.append(muscle_pair)
    all_muscles.append(shared_muscles)
    deformable.all_muscles = all_muscles

    # Implement the forward and backward simulation.
    dt = 3.33e-2
    num_frames = 200
    dofs = deformable.dofs()
    act_dofs = deformable.act_dofs()
    arrow_target_data = np.array([-1, 0, 0], dtype=np.float64)

    w_sideward = 10.0
    w_face = 0.0

    mid_x = math.floor(node_nums[0] / 2)
    mid_y = math.floor(node_nums[1] / 2)
    mid_z = math.floor(node_nums[2] / 2)
    mid_line = vertex_indices[:, mid_y, mid_z]
    center = vertex_indices[mid_x, mid_y, mid_z]

    face_head = vertex_indices[0, mid_y, mid_z]
    face_tail = vertex_indices[2, mid_y, mid_z]

    def get_state_(sim, q_, v_, a_=None, f_ext_=None):
        q_center = q_.reshape((-1, 3))[center]
        v_center = v_.reshape((-1, 3))[center]

        q_mid_line_rel = q_.reshape((-1, 3))[mid_line] - q_center
        v_mid_line = v_.reshape((-1, 3))[mid_line]
        state = [
            v_center,
            q_mid_line_rel.ravel(),
            v_mid_line.ravel(),
        ]
        return np.concatenate(state).copy()

    def get_reward_(sim, q_, v_, a_=None, f_ext_=None):

        v_center = np.mean(v_.reshape((-1, 3))[mid_line], axis=0)
        face_dir = q_.reshape((-1, 3))[face_head] - q_.reshape((-1, 3))[face_tail]
        face_dir = face_dir / np.linalg.norm(face_dir)

        # forward loss
        forward_reward = np.dot(v_center, arrow_target_data)

        # sideward loss
        cross = np.cross(v_center, arrow_target_data)
        sideward_reward = -np.dot(cross, cross)

        # face loss
        face_reward = np.dot(face_dir, arrow_target_data)

        return forward_reward + w_sideward * sideward_reward + w_face * face_reward

    def get_done_(sim, q_, v_, a_, f_ext_):
        if sim.frame >= sim.num_frames:
            return True
        return False

    setattr(sim_class, 'get_state_', get_state_)
    setattr(sim_class, 'get_reward_', get_reward_)
    setattr(sim_class, 'get_done_', get_done_)

    sim = sim_class(
        deformable, mesh, center, dofs, act_dofs, method, dt, opt, num_frames)

    if sim_class is AdaSim:
        action_shape = (len(all_muscles),)
    elif sim_class is IndSim:
        muscle_dofs = 0
        for shared_muscles in all_muscles:
            muscle_dofs += len(shared_muscles[0][0])
        action_shape = (muscle_dofs,)
    else:
        raise ValueError('invalid simulation class')

    sim.set_action_space(action_shape)

    sim.observation_space.seed(seed + rank)
    sim.action_space.seed(seed + rank)

    return sim

def make_starfish_3d(sim_class, seed, rank, *args, **kwargs):

    os.system("taskset -p 0xff %d" % os.getpid())

    seed = seed + rank
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.set_default_dtype(torch.float64)

    folder = Path('starfish').resolve()
    folder.mkdir(parents=True, exist_ok=True)

    # Mesh parameters
    limb_width = 2
    limb_length = 10
    limb_depth = 2

    cell_nums = [limb_length * 2 + limb_width, limb_length * 2 + limb_width, limb_depth]
    node_nums = [c + 1 for c in cell_nums]
    dx = 0.1
    origin = np.zeros((3,))
    bin_file_name = str(folder / 'starfish.bin')

    voxels = np.ones(cell_nums)
    voxels[:limb_length, :limb_length] = 0
    voxels[:limb_length, -limb_length:] = 0
    voxels[-limb_length:, :limb_length] = 0
    voxels[-limb_length:, -limb_length:] = 0

    voxel_indices, vertex_indices = generate_hex_mesh(
        voxels, dx, origin, bin_file_name, write=False)
    mesh = HexMesh3d()
    mesh.Initialize(bin_file_name)

    # FEM parameters.
    youngs_modulus = 1e6
    poissons_ratio = 0.45
    density = 1e3
    method = 'pd_eigen'
    opt = {
        'max_pd_iter': 1000, 'max_ls_iter': 10, 'abs_tol': 1e-4, 'rel_tol': 1e-3, 'verbose': 0,
        'thread_ct': 1, 'use_bfgs': 1, 'bfgs_history_size': 10
    }

    deformable = HexDeformable()
    deformable.Initialize(bin_file_name, density, 'none', youngs_modulus, poissons_ratio)
    # Elasticity.
    deformable.AddPdEnergy('corotated', [youngs_modulus / (1 + poissons_ratio),], [])
    # Hydrodynamics parameters.
    rho = 1e3
    v_water = [0, 0, 0]   # Velocity of the water.
    # # Cd_points = (angle, coeff) pairs where angle is normalized to [0, 1].
    Cd_points = ndarray([[0.0, 0.05], [0.4, 0.05], [0.7, 1.85], [1.0, 2.05]])
    # # Ct_points = (angle, coeff) pairs where angle is normalized to [-1, 1].
    Ct_points = ndarray([[-1, -0.8], [-0.3, -0.5], [0.3, 0.1], [1, 2.5]])
    # The current Cd and Ct are similar to Figure 2 in SoftCon.
    # surface_faces is a list of (v0, v1) where v0 and v1 are the vertex indices of the two endpoints of a boundary edge.
    # The order of (v0, v1) is determined so that following all v0 -> v1 forms a ccw contour of the deformable body.
    surface_faces = get_boundary_face(mesh)
    deformable.AddStateForce(
        'hydrodynamics', np.concatenate(
            [[rho,], v_water, Cd_points.ravel(), Ct_points.ravel(), ndarray(surface_faces).ravel()]))

    # Add actuation.
    all_muscles = []
    muscle_pairs = []

    muscle_stiffness = 1e5

    for move in [range(limb_length - 1, -1, -1), range(-limb_length, 0)]:
        for fix in [limb_length, limb_length + limb_width - 1]:

            muscle_pair = []
            for depth in [0, limb_depth - 1]:
                indices = [int(voxel_indices[fix, m, depth]) for m in move]
                deformable.AddActuation(muscle_stiffness, [0.0, 1.0, 0.0], indices)
                muscle_pair.append(indices)
            muscle_pairs.append(muscle_pair)

            muscle_pair = []
            for depth in [0, limb_depth - 1]:
                indices = [int(voxel_indices[m, fix, depth]) for m in move]
                deformable.AddActuation(muscle_stiffness, [1.0, 0.0, 0.0], indices)
                muscle_pair.append(indices)
            muscle_pairs.append(muscle_pair)

    all_muscles = [
        [muscle_pairs[0], muscle_pairs[2]],
        [muscle_pairs[1], muscle_pairs[3]],
        [muscle_pairs[4], muscle_pairs[6]],
        [muscle_pairs[5], muscle_pairs[7]],
    ]
    deformable.all_muscles = all_muscles

    # Implement the forward and backward simulation.
    dt = 3.33e-2
    num_frames = 200
    dofs = deformable.dofs()
    act_dofs = deformable.act_dofs()
    arrow_target_data = np.array([0, 0, 1], dtype=np.float64)

    w_sideward = 1.0
    w_face = 0.0

    mid_x = math.floor(node_nums[0] / 2)
    mid_y = math.floor(node_nums[1] / 2)
    mid_z = math.floor(node_nums[2] / 2)
    center = vertex_indices[mid_x, mid_y, mid_z]

    face_head = vertex_indices[mid_x, mid_y, -1]
    face_tail = vertex_indices[mid_x, mid_y, 0]

    mid_plane = np.array([
        vertex_indices[mid_x, :limb_length, mid_z],
        vertex_indices[mid_x, -limb_length:, mid_z],
        vertex_indices[:limb_length, mid_y, mid_z],
        vertex_indices[-limb_length:, mid_y, mid_z],
    ]).ravel()

    def get_state_(sim, q_, v_, a_=None, f_ext_=None):
        q_center = q_.reshape((-1, 3))[center]
        v_center = v_.reshape((-1, 3))[center]

        q_mid_line_rel = q_.reshape((-1, 3))[mid_plane] - q_center
        v_mid_line = v_.reshape((-1, 3))[mid_plane]
        state = [
            v_center,
            q_mid_line_rel.ravel(),
            v_mid_line.ravel(),
        ]
        return np.concatenate(state).copy()

    def get_reward_(sim, q_, v_, a_=None, f_ext_=None):

        v_center = v_.reshape((-1, 3))[center]
        face_dir = q_.reshape((-1, 3))[face_head] - q_.reshape((-1, 3))[face_tail]
        face_dir = face_dir / np.linalg.norm(face_dir)

        # forward loss
        forward_reward = np.dot(v_center, arrow_target_data)

        # sideward loss
        cross = np.cross(v_center, arrow_target_data)
        sideward_reward = -np.dot(cross, cross)

        # face loss
        face_reward = np.dot(face_dir, arrow_target_data)

        return forward_reward + w_sideward * sideward_reward + w_face * face_reward

    def get_done_(sim, q_, v_, a_, f_ext_):
        if sim.frame >= sim.num_frames:
            return True
        return False

    setattr(sim_class, 'get_state_', get_state_)
    setattr(sim_class, 'get_reward_', get_reward_)
    setattr(sim_class, 'get_done_', get_done_)

    sim = sim_class(
        deformable, mesh, center, dofs, act_dofs, method, dt, opt, num_frames)

    if sim_class is AdaSim:
        action_shape = (len(all_muscles),)
    elif sim_class is IndSim:
        muscle_dofs = 0
        for shared_muscles in all_muscles:
            muscle_dofs += len(shared_muscles[0][0])
        action_shape = (muscle_dofs,)
    else:
        raise ValueError('invalid simulation class')

    sim.set_action_space(action_shape)

    sim.observation_space.seed(seed + rank)
    sim.action_space.seed(seed + rank)

    return sim

class AdaSim(gym.Env):

    """Custom Environment that follows gym interface"""
    metadata = {'render.modes': ['human']}

    def __init__(
            self, deformable, mesh, center,
            dofs, act_dofs, method, dt, option, num_frames
        ):

        super(AdaSim, self).__init__()
        self.deformable = deformable
        self.mesh = mesh
        self.center = center
        self.dofs = dofs
        self.act_dofs = act_dofs
        self.method = method
        self.dt = dt
        self.option = option
        self.num_frames = num_frames

        self.a_init = np.zeros(self.act_dofs)
        self.prev_a = None

        self.frame = 0

        self.q = None
        self.v = None

        if isinstance(mesh, Mesh2d):
            dim = 2
        elif isinstance(mesh, HexMesh3d):
            dim = 3
        else:
            raise ValueError(f'invlaid mesh type: {type(mesh)}')

        q0 = ndarray(self.mesh.py_vertices())
        q0_center = q0.reshape((-1, dim))[center]
        self.q0 = (q0.reshape((-1, dim)) - q0_center).ravel()
        self.v0 = np.zeros_like(q0, dtype=np.float64)
        self.f_ext = np.zeros_like(q0, dtype=np.float64)

        self.observation_space = spaces.Box(
            low=-10.0, high=10.0, shape=self.get_state_(self.q0, self.v0).shape, dtype=np.float64) # pylint: disable=no-member
        self.action_space = None
        self.reset()

    def set_action_space(self, action_shape):
        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=action_shape, dtype=np.float64)

    def get_action(self, action):

        if self.prev_a is None:
            prev_a = self.a_init.copy()
        else:
            prev_a = self.prev_a

        a = []
        pointer = 0

        for w, shared_muscles in zip(action, self.deformable.all_muscles):
            mu_pair = [0.5 * (np.abs(w) - w), 0.5 * (np.abs(w) + w)]
            for muscle_pair in shared_muscles:
                if len(muscle_pair) != 2:
                    raise ValueError('adaptive controller require paired muscles')
                for mu, muscle in zip(mu_pair, muscle_pair):
                    prev_a_cord = prev_a[pointer:pointer + len(muscle)]
                    pointer += len(muscle)
                    a_cord = np.concatenate([mu.reshape((1,)), prev_a_cord[:-1]])
                    a.append(a_cord)

        a = np.array(a).ravel()
        self.prev_a = a.copy()
        return 1 - a

    def step(self, a):

        self.frame += 1

        a = self.get_action(a)

        q_array = StdRealVector(self.q)
        v_array = StdRealVector(self.v)
        a_array = StdRealVector(a)
        f_ext_array = StdRealVector(self.f_ext)
        q_next_array = StdRealVector(self.dofs)
        v_next_array = StdRealVector(self.dofs)

        self.deformable.PyForward(
            self.method, q_array, v_array, a_array, f_ext_array,
            self.dt, self.option, q_next_array, v_next_array, StdIntVector(0))

        q = ndarray(q_next_array)
        v = ndarray(v_next_array)

        self.q, self.v = q.copy(), v.copy()

        state = self.get_state_(q, v, a, self.f_ext) # pylint: disable=no-member
        reward = self.get_reward_(q, v, a, self.f_ext) # pylint: disable=no-member
        done = self.get_done_(q, v, a, self.f_ext) # pylint: disable=no-member

        if done:
            self.reset()

        return state, reward, done, dict()

    def reset(self):
        self.frame = 0
        self.q = self.q0.copy()
        self.v = self.v0.copy()
        self.prev_a = None

        return self.get_state_(self.q, self.v) # pylint: disable=no-member

class IndSim(AdaSim):
    def get_action(self, action):

        a_shared_muscles = []
        a = []

        pointer = 0
        for shared_muscles in self.deformable.all_muscles:
            a_shared_muscles.append(action[pointer:pointer + len(shared_muscles[0][0])])

        for w, shared_muscles in zip(a_shared_muscles, self.deformable.all_muscles):
            mu_pair = [0.5 * (np.abs(w) - w), 0.5 * (np.abs(w) + w)]
            for muscle_pair in shared_muscles:
                if len(muscle_pair) != 2:
                    raise ValueError('adaptive controller require paired muscles')
                for mu, muscle in zip(mu_pair, muscle_pair):
                    a.append(mu)

        a = np.array(a).ravel()
        return 1 - a


#!/usr/bin/env python3
"""
SWC: Manipulate graph data in SWC format.

Includes read/write to disk.
"""
from typing import Dict, List, Tuple, Union
import math

import networkx as nx
import numpy as np

class NodeTypes:
    """
    Enum of types for nodes in SWC.

    range(8)
    """

    UNDEFINED = 0
    SOMA = 1
    AXON = 2
    DENDRITE = 3
    APICAL_DENDRITE = 4
    FORK_POINT = 5
    END_POINT = 6
    CUSTOM = 7

    @staticmethod
    def is_valid(cls, t: int) -> bool:
        """
        Determine if a node type is a valid SWC-spec type.

        Arguments:
            t (int): Type

        Returns:
            bool: True if valid

        """
        return t in [0, 1, 2, 3, 4, 5, 6, 7]

class NeuronMorphology:
    """
    A wrapper class for neuron morphologies.

    Contains a graph representation of the morphology in nx.DiGraph format.
    """

    def __init__(self, **kwargs):
        """
        Create a new NeuronMorphology.

        Arguments:
            source (NeuronMorphology): Optional. Source to copy from

        """
        if "source" in kwargs:
            if isinstance(kwargs["source"], NeuronMorphology):
                self._skeleton = kwargs["source"].get_graph()
            elif isinstance(kwargs["source"], nx.Graph):
                self._skeleton = kwargs["source"]
            else:
                raise ValueError(
                    "The `source` argument passed to the NeuronMorphology "
                    + "constructor must be a graph or a NeuronMorphology."
                    + "Type was {}.".format(type(kwargs["source"]))
                )
        else:
            self._skeleton = nx.DiGraph()

    def __len__(self):
        return len(self._skeleton)

    def get_graph(self, copy: bool = True) -> nx.DiGraph:
        """
        Return the underlying graph data structure.

        By default, this returns a copy, so that modifications to this graph do
        not affect the parent NeuronMorphology (which is expected behavior).
        However, by explicitly passing `copy=False`, you can request a pointer
        to the same graph.

        Arguments:
            copy (bool : True): If a copy should be returned instead of a
                pointer to the original

        Returns:
            networkx.Graph: The graph data structure that backs the morphology

        """
        if copy:
            return self._skeleton.copy()
        else:
            return self._skeleton

    def add_node(
        self, id: int, t: int = None, xyz: Tuple[int, int, int] = None, r: float = None
    ) -> None:
        """
        Add a new node to the skeleton.

        Arguments:
            id (int): The ID of the node (0 is invalid)
            t (int): The type of node. Validate with `NodeTypes.is_valid(t)`
            xyz (int[3]): The xyz position of the node
            r (float): The radius of the neuron at this node

        Returns:
            None
        """
        return self._skeleton.add_node(id, t=t, xyz=xyz, r=r)

    def add_edge(self, start: int, end: int) -> None:
        """
        Add a new edge to the skeleton.

        Arguments:
            start (int): The origin of the edge (ID)
            stop (int): The end of the edge (ID)

        Returns:
            None

        """
        return self._skeleton.add_edge(start, end)

    def get_branch_points(self) -> List[int]:
        """
        Returns a list of all node IDs where degree > 2.

        Arguments:
            None

        Returns:
            int[]: Node IDs where degree > 2
        """
        results: List[int] = []
        for start, stops in self._skeleton.pred.items():
            if len(stops.keys()) >= 2:
                results.append(start)
        return results

    def get_branch_angles(self) -> Dict[int, float]:
        """Get the branch angles between all connected nodes."""

        angles = {}
        for node in self._skeleton.nodes():
            neighbors = [i for i in nx.all_neighbors(self._skeleton, node)]
            if (len(neighbors)) == 2:
                # This is not a branch point, so we can compute:
                angles[node] = self.get_branch_angle((neighbors[0], node, neighbors[1]))
        return angles

    def get_branch_angle(self, abc: Tuple[int, int, int]):
        """
        Returns the minimum branch angle between edges AB and BC.
        """
        a, b, c = abc
        apos = np.array(self._skeleton.nodes[a]["xyz"])
        bpos = np.array(self._skeleton.nodes[b]["xyz"])
        cpos = np.array(self._skeleton.nodes[c]["xyz"])

        ba_unit_vector = apos - bpos
        ba_unit_vector /= np.linalg.norm(ba_unit_vector)
        bc_unit_vector = cpos - bpos
        bc_unit_vector /= np.linalg.norm(bc_unit_vector)
        dot_product = np.dot(ba_unit_vector, bc_unit_vector)
        return np.arccos(dot_product)

    def get_distance_between_nodes(self, a: int, b: int) -> float:
        """
        Returns the distance between two nodes.

        Arguments:
            a (int): The first node ID
            b (int): The second node ID

        Returns:
            float: The distance between the two nodes
        """
        return np.linalg.norm(
            np.array(self._skeleton.nodes[a]["xyz"])
            - np.array(self._skeleton.nodes[b]["xyz"])
        )

    def get_path_length(self, start: int, end: int) -> float:
        """
        Get the path length between two nodes.

        Arguments:
            start (int): The origin of the path (ID)
            stop (int): The end of the path (ID)

        Returns:
            float: The length of the path

        """
        # Get the shortest path IDs:
        path = nx.shortest_path(self._skeleton, start, end)
        # Compute the path length. The length is the sum of x1y1z1 - x0y0z0:
        length = 0
        for i in range(len(path) - 1):
            length += self.get_distance_between_nodes(path[i], path[i + 1])
        return length

    def get_total_length(self) -> float:
        """
        Returns the total length of the neuron's segments.
        """
        # Iterate over all edges and add the length of each edge:
        total_length = 0
        for start, end in self._skeleton.edges():
            total_length += self.get_distance_between_nodes(start, end)
        return total_length

    def smoothed(self) -> nx.DiGraph:
        """
        Returns a _copy_ of this morphology as a smoothed graph.
        # TODO: This is very inefficient.

        """
        gcopy = self.get_graph()
        gcopy_old = self.get_graph()
        new_count = 0
        old_count = len(self._skeleton.adj)
        while new_count != old_count:
            gcopy_old = gcopy
            old_count = len(gcopy_old.adj)
            for node, connections in gcopy_old.adj.items():
                if len(connections) == 2:
                    start, stop = connections.keys()
                    gcopy.add_edge(start, stop)
                    gcopy.remove_node(node)
                    break
            new_count = len(gcopy.adj)
        return gcopy

    @staticmethod
    def from_file(filename: str):
        return load_swc(filename)

    @staticmethod
    def from_string(swc: str):
        return read_swc(swc)

    def translate(self, translation: Tuple[int, int, int], inplace=False):
        """
        Translate the target neuron morphology (affine translation) in XYZ.

        Arguments:
            translation (Tuple[int, int, int]): The translation to perform
            inplace (bool: False): Whether to perform the translation on this
                morphology (True) or on a copy (False).

        Returns:
            The morphology upon which the translation was performed

        """
        if inplace:
            target = self
        else:
            target = NeuronMorphology(source=self)
        for node in target._skeleton.nodes():
            current = target._skeleton.nodes[node]["xyz"]
            target._skeleton.nodes[node]["xyz"] = [
                current[0] + translation[0],
                current[1] + translation[1],
                current[2] + translation[2],
            ]
        return target

    def scale(self, scale: Union[float, Tuple[float, float, float]], inplace=False):
        """
        Scale the target neuron morphology.

        Arguments:
            scale (Union[float, Tuple[float, float, float]]): The scale to
                perform. If a tuple, [X,Y,Z]. If a scalar, perform an isometric
                scale on all three axes.
            inplace (bool: False): Whether to perform the translation on this
                morphology (True) or on a copy (False).

        Returns:
            The morphology upon which the scaling was performed

        """
        if inplace:
            target = self
        else:
            target = NeuronMorphology(source=self)

        if isinstance(scale, (float, int)):
            scale = (scale, scale, scale)

        for node in target._skeleton.nodes():
            current = target._skeleton.nodes[node]["xyz"]
            target._skeleton.nodes[node]["xyz"] = [
                current[0] * scale[0],
                current[1] * scale[1],
                current[2] * scale[2],
            ]
        return target

    def rotate(
        self, rotation: Tuple[int, int, int], inplace: bool = False, _p: int = 10
    ):
        """
        Perform a rotation on the neuron morphology.

        If inplace is True, will perform the rotation on the current object.

        Arguments:
            rotation (Tuple[int, int, int]): The rotation to perform, in
                pitch-roll-yaw order. Units are in radians.
            inplace (bool: False): Whether to perform the rotation on this
                morphology (True) or on a copy (False).
            _p (int: 10): Digits of precision. Because SWC files are text-
                files, higher precision here costs more in terms of file
                storage. Values of 5 or greater should be fine for most
                purposes, unless coordinates are small to begin with.

        Returns:
            The morphology upon which the translation was performed

        """

        # Formula adapted from https://stackoverflow.com/a/34060479/979255
        pitch, roll, yaw = rotation
        cosa = math.cos(yaw)
        sina = math.sin(yaw)

        cosb = math.cos(pitch)
        sinb = math.sin(pitch)

        cosc = math.cos(roll)
        sinc = math.sin(roll)

        Axx = cosa * cosb
        Axy = cosa * sinb * sinc - sina * cosc
        Axz = cosa * sinb * cosc + sina * sinc

        Ayx = sina * cosb
        Ayy = sina * sinb * sinc + cosa * cosc
        Ayz = sina * sinb * cosc - cosa * sinc

        Azx = -sinb
        Azy = cosb * sinc
        Azz = cosb * cosc

        if inplace:
            target = self
        else:
            target = NeuronMorphology(source=self)

        for node in target._skeleton.nodes():
            current = target._skeleton.nodes[node]["xyz"]
            target._skeleton.nodes[node]["xyz"] = [
                round(Axx * current[0] + Axy * current[1] + Axz * current[2], _p),
                round(Ayx * current[0] + Ayy * current[1] + Ayz * current[2], _p),
                round(Azx * current[0] + Azy * current[1] + Azz * current[2], _p),
            ]
        return target

    def draw(self, node_radius_multiplier: int = 10):
        k = self._skeleton
        pos = {n: a["xyz"][:2] for n, a in k.nodes(data=True)}
        nx.draw(
            nx.Graph(k),
            pos=pos,
            width=5,
            node_size=[n["r"] * node_radius_multiplier for _, n in k.nodes(data=True)],
            arrowsize=1,
        )

def read_swc(swc_str: str) -> NeuronMorphology:
    """
    Construct a NeuronMorphology from a SWC string.

    For file imports, see also `load_swc`.

    Returns:
        NeuronMorphology
    """
    lines = swc_str.split("\n")
    neuron = NeuronMorphology()
    last_index = None
    for line in lines:
        line = line.strip()
        if (not line) or (line[0] == "#"):
            continue
        else:
            attrs = [float(i) for i in line.split()]
            neuron.add_node(int(attrs[0]), t=int(attrs[1]), xyz=attrs[2:5], r=attrs[5])
            last_index = attrs[-1]
            if last_index >= 0:
                neuron.add_edge(int(attrs[0]), int(attrs[-1]))
    return neuron

def load_swc(filename: str) -> NeuronMorphology:
    """
    Loads a SWC from disk, into a NeuronMorphology object.

    For str imports, see also `read_swc`.

    Arguments:
        filename (str)

    Returns:
        NeuronMorphology

    """
    try:
        with open(filename, "r") as fh:
            contents = fh.read()
            return read_swc(contents)
    except Exception as ex:
        raise ValueError("Invalid file {}".format(filename))

def save_swc(filename: str, nmorpho: NeuronMorphology) -> str:
    """
    Saves a morphology to disk in the form of a SWC file.

    Arguments:
        filename (str): The file to which to save the SWC
        nmorpho (NeuronMorphology): The morphology to save

    Returns:
        str: File path on disk to which the SWC was saved

    """
    lines = []
    g = nmorpho.get_graph()
    _edges = g.edges()
    # Loop through the nodes. Pass `True` to include metadata:
    for node in g.nodes(data=True):
        parent = [i for i in g.succ[node[0]]]
        if parent == []:
            parent = -1
        else:
            parent = parent[0]

        #             n  T xyz R  P
        lines.append(
            "{} {} {} {} {}".format(
                str(node[0]),
                str(node[1]["t"]),
                " ".join([str(i) for i in node[1]["xyz"]]),
                str(node[1]["r"]),
                str(parent),
            )
        )
    with open(filename, "w") as swc_output:
        swc_output.write("\n".join(lines))
        swc_output.write("\n")
    return filename


##
## function to return the AGN spectrum
##   given nu (in Hz), and log(L_bol/L_sun), returns
##   log(nuL_nu/L_sun)
##
##	keywords to shortcut for various bands:
##	  BB = B-band, IR = mid-IR (15 microns),
##    SX = soft X-rays (0.5-2 keV),
##    HX = hard X-rays (2-10 keV)
##
##  keywords for different spectra:
##    MARCONI = Marconi et al. 2004-ish spectrum -- it's not *really*
##      the same (updated w. more recent estimates of the X-ray
##      input parameters, i.e. Gamma & reflection angle, from Tozzi et al. 2006
##		as well as the more recent Steffen et al. 2006 calibration of the
##      alpha_ox relation (specifically their l_uv-l_2keV bisector measurement,
##      which is a subtle difference but is very important for the bolometric
##      corrections of the most luminous X-ray sources
##    HRH = Hopkins, Richards, & Hernquist 2006 template spectrum
##	    compiled from a number of observations therein. It's a little closer to
##      a typical observed spectrum, with a more detailed modeling of a number of
##      parts of the spectrum (i.e. more continuum shape). the major difference is
##      that it includes a hot dust component, which is important to account for
##      since much of the energy comes out there. however, if you want to model
##      that generation self-consistently, you should use the "input" marconi et al.
##      spectrum. still, the galaxy-scale obscuration doesn't necessarily produce this
##      and the alpha-ox calibrations are for objects with this feature, so it is
##      generally more representative
##    RICHARDS = Richards et al. 2006 all-quasar mean SED: if you want a
##      non-luminosity-dependent spectrum for whatever reason, this is an improved
##      version of Elvis et al. 1994 (which is too X-ray bright)
##
##    SDSS can be added as a keyword with any of the model spectra, and it will
##      overlay the vanden Berk et al. 2001 median SDSS SED over the relevant
##      portion of the spectrum :: basically does so with a sliding continuum determination
##      that then overlays this spectrum, such that the integrated luminosity in the
##      entire bandpass and the continuum spectrum are conserved (i.e. you can put this
##      over an arbitrary continuum
##
##	unless you're using the Richards et al. 2006 spectra, if you want a reference for this,
##    it should be Hopkins, Richards, & Hernquist 2006. Even the "MARCONI" key spectrum
##    is substantially modified as described. If you want the additional relevant
##    observational compilations on which it's all based, the list is :
##    Richards et al. 2006, Steffen et al. 2006, Hatziminaoglou et al. 2005,
##    Tozzi et al. 2006, Strateva et al. 2005, Telfer et al. 2002, Vanden Berk et al. 2001,
##    George et al. 1998, Elvis et al. 1994, Marconi et al. 2004 (not really an observational
##    paper but the methodology does follow them), with the X-ray reflection component
##    following Ueda et al. 2003 in the PEXRAV code, Magdziarz & Zdziarski 1995
##
##
##
def agn_spectrum( nu_in_Hz, log_l_bol, \
	BB=0, IR=0, SX=0, HX=0, \
	HRH=0, MARCONI=0, RICHARDS=0, \
	SDSS=0 ):

    import utilities as util
    import numpy as np
    import math
    import ctypes

    ## location of shared library
    exec_call=util.return_python_routines_homedir()+'/agn_spectrum/agn_spectrum_py.so'
    lib=ctypes.cdll[exec_call];

    if (1 == BB) : nu_in_Hz = -1.0
    if (1 == IR) : nu_in_Hz = -2.0
    if (1 == SX) : nu_in_Hz = -3.0
    if (1 == HX) : nu_in_Hz = -4.0
    spectrum_key = 0
    if (1 == HRH) 		: spectrum_key = 0
    if (1 == MARCONI) 	: spectrum_key = 1
    if (1 == RICHARDS) 	: spectrum_key = 2
    sloan_key    = 0
    if (1 == SDSS)		: sloan_key = 1

    nu_in_Hz=np.array(nu_in_Hz,ndmin=1,dtype='d');
    log_l_bol=np.array(log_l_bol,ndmin=1,dtype='d');

    N_nu = nu_in_Hz.shape[0]
    N_lum = log_l_bol.shape[0]
    l_band_all = np.zeros((N_nu,N_lum),dtype='d')

    out_cast = ctypes.c_double*N_nu
    for i in range(N_lum):
        #l_band_vec = np.zeros(N_nu,dtype='d')
        l_band_vec = out_cast()
        log_l_bol_pass = log_l_bol[i]

        lib.main( ctypes.c_int(N_nu),\
            nu_in_Hz.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),\
            ctypes.c_double(log_l_bol_pass),\
            ctypes.c_int(spectrum_key),\
            ctypes.c_int(sloan_key),\
            ctypes.byref(l_band_vec) );
            #l_band_vec.ctypes.data_as(ctypes.POINTER(ctypes.c_double)));
        ## now put the output arrays into a useful format
        l_band_vec = np.copy(np.ctypeslib.as_array(l_band_vec));
        l_band_all[:,i] = l_band_vec

    return l_band_all


# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# 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.

"""Mixture of Experts routing mechanisms."""

from typing import Any, Callable, Iterable, Tuple

import flax
from flax import linen as nn
import jax
import jax.numpy as jnp

# Type Stubs
PRNGKey = Any
RouterOutput = Any
Shape = Iterable[int]

# Switch Transformer (https://arxiv.org/abs/2101.03961) suggests using
# nn.initializers.variance_scaling(0.1, "fan_in", "truncated_normal")
# scaling throughout MoE models, but we find slightly better results adopting
# typical normally-distributed scaling for the router specifically.
default_kernel_init = nn.initializers.normal(stddev=2e-2)
default_bias_init = nn.initializers.zeros

@flax.struct.dataclass
class RouterIndices:
  """Dispatch indices and combine weights for scatter/gather-based routing.

  Attributes:
    dispatch_indices: <int32>[NUM_GROUPS, TOKENS_PER_GROUP,
      NUM_SELECTED_EXPERTS, 2] dispatch indices indicating, for each token, its
      preferred expert and its priority in that expert's buffer.
    combine_weights: <float>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS]
      combine weights used for scaling expert outputs with the router's dispatch
      probability/confidence.
    auxiliary_loss: Load balancing loss for router.
    router_z_loss: Router z-loss. Encourages router logits to remain small in an
      effort to improve stability.
  """
  dispatch_indices: jnp.ndarray
  combine_weights: jnp.ndarray
  auxiliary_loss: float
  router_z_loss: float = 0.

@flax.struct.dataclass
class RouterMask:
  """Dispatch and combine arrays for expert routing with masked matmuls.

  Attributes:
    dispatch_mask: <bool>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS,
      EXPERT_CAPACITY] dispatch array that is 1 if the token gets routed to the
      corresponding expert, and 0 otherwise.
    combine_array: <float>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS,
      EXPERT_CAPACITY] combine array used for combining expert outputs and
      scaling with router probability.
    auxiliary_loss: Load balancing loss for router.
    router_z_loss: Router z-loss. Encourages router logits to remain small in an
      effort to improve stability.
  """
  dispatch_mask: jnp.ndarray
  combine_array: jnp.ndarray
  auxiliary_loss: float
  router_z_loss: float = 0.

def _favor_one_hot_slices():
  """Returns true iff running on TPUs."""
  return jax.default_backend() == "tpu" or jax.devices()[0].platform == "tpu"

def _take_along_axis(array, indices,
                     axis):
  """Takes values from the input array by matching 1D index and data slices.

  This function serves the same purpose as jax.numpy.take_along_axis, except
  that it uses one-hot matrix multiplications under the hood on TPUs:
  (1) On TPUs, we use one-hot matrix multiplications to select elements from the
      array.
  (2) Otherwise, we fall back to jax.numpy.take_along_axis.

  Notes:
    - To simplify matters in case (1), we only support slices along the second
      or last dimensions.
    - We may wish to revisit (1) for very large arrays.

  Args:
    array: Source array.
    indices: Indices to take along each 1D slice of array.
    axis: Axis along which to take 1D slices.

  Returns:
    The indexed result.
  """
  if array.ndim != indices.ndim:
    raise ValueError(
        "indices and array must have the same number of dimensions; "
        f"{indices.ndim} vs. {array.ndim}.")

  if (axis != -1 and axis != array.ndim - 1 and  # Not last dimension
      axis != 1 and axis != -array.ndim + 1):  # Not second dimension
    raise ValueError(
        "Only slices along the second or last dimension are supported; "
        f"array.ndim = {array.ndim}, while axis = {axis}.")

  if _favor_one_hot_slices():
    one_hot_length = array.shape[axis]
    one_hot_indices = jax.nn.one_hot(indices, one_hot_length, axis=axis)

    if axis == -1 or array.ndim == 1:
      # Take i elements from last dimension (s).
      # We must use HIGHEST precision to accurately reproduce indexing
      # operations with matrix multiplications.
      result = jnp.einsum(
          "...s,...is->...i",
          array,
          one_hot_indices,
          precision=jax.lax.Precision.HIGHEST)
    else:
      # Take i elements from second dimension (s). We assume here that we always
      # want to slice along the second dimension.
      # We must use HIGHEST precision to accurately reproduce indexing
      # operations with matrix multiplications.
      result = jnp.einsum(
          "ns...,nis...->ni...",
          array,
          one_hot_indices,
          precision=jax.lax.Precision.HIGHEST)
    return jax.lax.convert_element_type(result, array.dtype)
  else:
    return jnp.take_along_axis(array, indices, axis=axis)

def _top_k(array, k):
  """Returns top k values and their indices along the last axis of the array.

  This function serves the same purpose as jax.lax.top_k, but in a more XLA
  friendly manner for TPUs:
  (1) On TPUs, we use one-hot matrix multiplications to select the top k values.
      This convoluted way of obtaining the top k values is generally faster on
      TPUs.
  (2) Otherwise, we fall back to jax.lax.top_k (and its underlying scatter op).

  Args:
    array: Source array.
    k: Number of top values to select.

  Returns:
    - Top k values
    - Associated top k indices.
  """
  if _favor_one_hot_slices():
    top_k_indices = jax.lax.top_k(array, k)[-1]
    top_k_values = _take_along_axis(array, top_k_indices, axis=-1)
    return top_k_values, top_k_indices
  else:
    return jax.lax.top_k(array, k)

class RouterWeights(nn.Module):
  """Router module converting token inputs to router logits.

  Attributes:
    use_bias: Whether or not to use the bias term in computing the logits.
    dtype: Numerical float type for router logit computation.
    kernel_init: Initialization scheme for kernel.
    bias_init: Initialization scheme for bias.
    precision: XLA precision for array computations.
  """
  use_bias: bool = True
  dtype: jnp.dtype = jnp.bfloat16
  kernel_init: Callable[[PRNGKey, Shape, jnp.dtype],
                        jnp.ndarray] = default_kernel_init
  bias_init: Callable[[PRNGKey, Shape, jnp.dtype],
                      jnp.ndarray] = default_bias_init
  precision: jax.lax.Precision = jax.lax.Precision.DEFAULT

  @nn.compact
  def __call__(self, token_inputs,
               num_experts):
    """Applies RouterWeights module.

    Args:
      token_inputs: Flattened batch of tokens with shape <float>[NUM_GROUPS,
        TOKENS_PER_GROUP, HIDDEN_DIM].
      num_experts: Number of experts.

    Returns:
      Router logits with shape <float>[NUM_GROUPS, TOKENS_PER_GROUP,
      NUM_EXPERTS].
    """
    return nn.DenseGeneral(
        num_experts,
        use_bias=self.use_bias,
        dtype=self.dtype,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        precision=self.precision)(
            token_inputs)

class Router(nn.Module):
  """Abstract base router class, defining router API and inner workings.

  Attributes:
    router_weights: Configurable module used to compute router logits from token
      inputs.
    jitter_noise: Amplitude of jitter noise applied to router logits.
    dtype: Numeric float type for returned combine array. All actual
      computations are performed in float32 of the input for stability.
  """
  router_weights: RouterWeights
  jitter_noise: float
  dtype: jnp.dtype

  def __call__(self,
               token_inputs,
               num_experts,
               expert_capacity,
               apply_jitter = True):
    """Computes dispatch and combine arrays for routing to experts.

    Args:
      token_inputs: <float>[NUM_GROUPS, TOKENS_PER_GROUP, HIDDEN_DIM] inputs to
        send to experts.
      num_experts: Number of experts.
      expert_capacity: Each group will send this many tokens to each expert.
      apply_jitter: If true, apply jitter noise during routing.

    Returns:
      Router indices or mask arrays (depending on router type).
    """
    router_probs, router_logits = self._compute_router_probabilities(
        token_inputs, num_experts, apply_jitter)
    instructions = self._compute_routing_instructions(router_probs,
                                                      expert_capacity)
    return instructions.replace(router_z_loss=_router_z_loss(router_logits))

  def _compute_router_probabilities(
      self, token_inputs, num_experts,
      apply_jitter):
    """Computes router probabilities from input tokens.

    Args:
      token_inputs: <float>[NUM_GROUPS, TOKENS_PER_GROUP, HIDDEN_DIM] from which
        router probabilities are computed.
      num_experts: Number of experts.
      apply_jitter: If true, apply jitter noise.

    Returns:
      - <float32>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS] probabilities for
        each token and expert. Used for routing tokens to experts.
      - <float>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS] raw router logits.
        Used for computing router z-loss.
    """
    # For remainder of routing computation we use float32 to ensure stability.
    # See the discussion of "selective precision" in
    # https://arxiv.org/abs/2101.03961.
    token_inputs = jax.lax.convert_element_type(token_inputs, jnp.float32)

    if apply_jitter and self.jitter_noise > 0:
      token_inputs *= jax.random.uniform(
          self.make_rng("jitter"),
          token_inputs.shape,
          token_inputs.dtype,
          minval=1.0 - self.jitter_noise,
          maxval=1.0 + self.jitter_noise)

    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS]
    router_logits = self.router_weights(token_inputs, num_experts)

    router_probabilities = jax.nn.softmax(router_logits, axis=-1)

    return router_probabilities, router_logits

  def _compute_routing_instructions(self, router_probs,
                                    expert_capacity):
    """Computes instructions for routing inputs to experts."""
    raise NotImplementedError(
        "Router is an abstract class that should be subclassed.")

class ScatterRouter(Router):
  """Abstract base router class for scatter dispatch routers.

  ScatterRouter(s) return RouterIndices containing dispatch indices and combine
  weights for sending token inputs (via scatter) and receiving outputs (via
  gather) to and from experts.

  Scatter-based routing is generally faster than masked matmul routing on CPUs
  and GPUs.
  """

  def _compute_routing_instructions(self, router_probs,
                                    expert_capacity):
    """Computes instructions for routing inputs to experts.

    Args:
      router_probs: <float32>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS]
        probabilities used to determine the routing of tokens to the experts.
      expert_capacity: Each group will send this many tokens to each expert.

    Returns:
      Router indices containing dispatch indices and combine weights.
    """
    raise NotImplementedError(
        "ScatterRouter is an abstract class that should be subclassed.")

class MaskedRouter(Router):
  """Abstract base router class for masked matmul dispatch routers.

  MaskedRouter(s) return RouterMask(s) containing a dispatch mask and combine
  array for sending and receiving (via masked matmuls) inputs and outputs to and
  from experts.

  Routing using masked matmuls is generally faster than scatter-based routing on
  CPUs and GPUs.
  """

  def _compute_routing_instructions(self, router_probs,
                                    expert_capacity):
    """Computes masks for the top-k experts per token.

    Args:
      router_probs: <float32>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS]
        probabilities used to determine the routing of tokens to the experts.
      expert_capacity: Each group will send this many tokens to each expert.

    Returns:
      Router mask arrays.
    """
    raise NotImplementedError(
        "MaskedRouter is an abstract class that should be subclassed.")

class TokensChooseScatterRouter(ScatterRouter):
  """Scatter router using tokens choose top-k experts assignment.

  This router uses the same mechanism as in Switch Transformer
  (https://arxiv.org/abs/2101.03961): tokens choose their top experts. Items are
  sorted by router_probs and then routed to their choice of expert until the
  expert's expert_capacity is reached. There is no guarantee that each token is
  processed by an expert, or that each expert receives at least one token.

  Attributes:
    num_selected_experts: Maximum number of experts to which each token is
      routed. Tokens may be routed to fewer experts if particular experts are
      oversubscribed / reach capacity.
    batch_prioritized_routing: Whether or not to use Batch Prioritized Routing
      (BPR), originally introduced in V-MoE (https://arxiv.org/abs/2106.05974).
        With BPR, we prioritize routing those top-k tokens with the highest
        router probability, rather than simply using each tokens left-to-right
        ordering in the batch. This prioritization is important because the
        experts have limited capacity.
  """
  num_selected_experts: int
  batch_prioritized_routing: bool

  def _compute_routing_instructions(self, router_probs,
                                    expert_capacity):
    """Computes dispatch indices and combine weights for the top-k experts.

    Args:
      router_probs: <float32>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS]
        probabilities used to determine the routing of tokens to the experts.
      expert_capacity: Each group will send this many tokens to each expert.

    Returns:
      Dispatch indices and combine weights for scatter/gather-based routing.
    """
    num_groups, tokens_per_group, num_experts = router_probs.shape

    # Top-k router probability and corresponding expert indices for each token.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS].
    combine_weights, expert_indices = _top_k(
        router_probs, k=self.num_selected_experts)

    auxiliary_loss = _load_balancing_loss(router_probs, expert_indices)

    if self.batch_prioritized_routing:
      # Sort tokens according to their routing probability per token group, so
      # that the highest probability tokens are routed first.
      token_ordering = jnp.argsort(-combine_weights[Ellipsis, 0], axis=-1)
      expert_indices = _take_along_axis(
          expert_indices, jnp.expand_dims(token_ordering, axis=-1), axis=-2)

    # Identify each token's preferred expert.
    # Make NUM_SELECTED_EXPERTS the leading axis to ensure that top-1 choices
    # have priority over top-2 choices, which have priority over top-3
    # choices...
    preferred_experts = jnp.swapaxes(expert_indices, 1, 2)
    # Shape: [NUM_GROUPS, NUM_SELECTED_EXPERTS * TOKENS_PER_GROUP]
    preferred_experts = preferred_experts.reshape(num_groups, -1)

    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP * NUM_SELECTED_EXPERTS, NUM_EXPERTS].
    expert_mask = jax.nn.one_hot(
        preferred_experts, num_experts, dtype=jnp.int32)

    # Experts have a fixed capacity that we cannot exceed. A token's priority
    # within the expert's buffer is given by the masked, cumulative capacity of
    # its target expert.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP * NUM_SELECTED_EXPERTS, NUM_EXPERTS].
    token_priority = jnp.cumsum(expert_mask, axis=1) * expert_mask - 1.0
    # Shape: [NUM_GROUPS, NUM_SELECTED_EXPERTS, TOKENS_PER_GROUP, NUM_EXPERTS].
    token_priority = token_priority.reshape(
        (num_groups, self.num_selected_experts, -1, num_experts))
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS, NUM_EXPERTS].
    token_priority = jnp.swapaxes(token_priority, 1, 2)
    # For each token, across all experts, select the only non-negative
    # (unmasked) priority.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS].
    token_priority = jnp.max(token_priority, axis=-1)

    # Return to original index shape.
    preferred_experts = preferred_experts.reshape(num_groups,
                                                  self.num_selected_experts,
                                                  tokens_per_group)
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS]
    preferred_experts = jnp.swapaxes(preferred_experts, 1, 2)

    if self.batch_prioritized_routing:
      # Place tokens in their original ordering.
      inverse_token_ordering = jnp.argsort(token_ordering, axis=-1)
      preferred_experts = _take_along_axis(
          preferred_experts,
          jnp.expand_dims(inverse_token_ordering, axis=-1),
          axis=-2)
      token_priority = _take_along_axis(
          token_priority,
          jnp.expand_dims(inverse_token_ordering, axis=-1),
          axis=-2)

    # Mask out tokens that overflow the maximum expert capacities.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS].
    combine_weights *= token_priority < expert_capacity

    # Expert index and priority within the expert capacity buffer.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS, 2].
    dispatch_indices = jnp.stack([preferred_experts, token_priority], axis=-1)

    # Return to default dtype now that router computation is complete.
    combine_weights = jax.lax.convert_element_type(combine_weights, self.dtype)
    dispatch_indices = jax.lax.convert_element_type(dispatch_indices, jnp.int32)

    return RouterIndices(dispatch_indices, combine_weights, auxiliary_loss)

class TokensChooseMaskedRouter(MaskedRouter):
  """Masked matmul router using tokens choose top-k experts assignment.

  This router uses the same mechanism as in Switch Transformer
  (https://arxiv.org/abs/2101.03961): tokens choose their top experts. Items are
  sorted by router_probs and then routed to their choice of expert until the
  expert's expert_capacity is reached. There is no guarantee that each token is
  processed by an expert, or that each expert receives at least one token.

  Attributes:
    num_selected_experts: Maximum number of experts to which each token is
      routed. Tokens may be routed to fewer experts if particular experts are
      oversubscribed / reach capacity.
    batch_prioritized_routing: Whether or not to use Batch Prioritized Routing
      (BPR), originally introduced in V-MoE (https://arxiv.org/abs/2106.05974).
        With BPR, we prioritize routing those top-k tokens with the highest
        router probability, rather than simply using each tokens left-to-right
        ordering in the batch. This prioritization is important because the
        expert's have limited capacity.
  """
  num_selected_experts: int
  batch_prioritized_routing: bool

  def _compute_routing_instructions(self, router_probs,
                                    expert_capacity):
    """Computes masks for the top-k experts per token.

    Args:
      router_probs: <float32>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS]
        probabilities used to determine the routing of tokens to the experts.
      expert_capacity: Each group will send this many tokens to each expert.

    Returns:
      Dispatch and combine arrays for routing with masked matmuls.
    """
    num_groups, _, num_experts = router_probs.shape

    # Top-k router probability and corresponding expert indices for each token.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS].
    expert_gate, expert_index = _top_k(
        router_probs, k=self.num_selected_experts)

    auxiliary_loss = _load_balancing_loss(router_probs, expert_index)

    if self.batch_prioritized_routing:
      # Sort tokens according to their routing probability per group, so that
      # the highest probability tokens are routed first.
      permutation = jnp.argsort(-expert_gate[Ellipsis, 0], axis=-1)
      # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS]
      expert_index = _take_along_axis(
          expert_index, jnp.expand_dims(permutation, axis=-1), axis=-2)

    # Make NUM_SELECTED_EXPERTS the leading axis to ensure that top-1 choices
    # have priority over top-2 choices, which have priority over top-3 choices,
    # etc.
    expert_index = jnp.swapaxes(expert_index, 1, 2)
    # Shape: [NUM_GROUPS, NUM_SELECTED_EXPERTS * TOKENS_PER_GROUP]
    expert_index = expert_index.reshape(num_groups, -1)

    # Create mask out of indices.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP * NUM_SELECTED_EXPERTS, NUM_EXPERTS].
    expert_mask = jax.nn.one_hot(expert_index, num_experts, dtype=jnp.int32)

    # Experts have a fixed capacity that we cannot exceed. A token's priority
    # within the expert's buffer is given by the masked, cumulative capacity of
    # its target expert.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP * NUM_SELECTED_EXPERTS, NUM_EXPERTS].
    token_priority = jnp.cumsum(expert_mask, axis=1) * expert_mask - 1.0
    # Shape: [NUM_GROUPS, NUM_SELECTED_EXPERTS, TOKENS_PER_GROUP, NUM_EXPERTS].
    token_priority = token_priority.reshape(
        (num_groups, self.num_selected_experts, -1, num_experts))
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS, NUM_EXPERTS].
    token_priority = jnp.swapaxes(token_priority, 1, 2)
    # For each token, across all selected experts, select the only non-negative
    # (unmasked) priority. Now, for group G routing to expert E, token T has
    # non-negative priority (i.e. token_priority[G,T,E] >= 0) if and only if E
    # is its targeted expert.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS].
    token_priority = jnp.max(token_priority, axis=2)

    if self.batch_prioritized_routing:
      # Place token priorities in original ordering of tokens.
      inv_permutation = jnp.argsort(permutation, axis=-1)
      token_priority = _take_along_axis(
          token_priority, jnp.expand_dims(inv_permutation, axis=-1), axis=-2)

    # Token T can only be routed to expert E if its priority is positive and
    # less than the expert capacity. One-hot matrix will ignore indices outside
    # the range [0, EXPERT_CAPACITY).
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS, EXPERT_CAPACITY].
    dispatch_mask = jax.nn.one_hot(
        token_priority, expert_capacity, dtype=jnp.bool_)

    # The combine array will be used for combining expert outputs, scaled by the
    # router probabilities. Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS,
    # EXPERT_CAPACITY].
    combine_array = jnp.einsum(
        "...te,...tec->...tec",
        router_probs,
        dispatch_mask,
        precision=jax.lax.Precision.DEFAULT)

    # Return to default dtype now that router computation is complete.
    combine_array = jax.lax.convert_element_type(combine_array, self.dtype)

    return RouterMask(dispatch_mask, combine_array, auxiliary_loss)

class ExpertsChooseMaskedRouter(MaskedRouter):
  """Masked matmul router using experts choose tokens assignment.

  This router uses the same mechanism as in Mixture-of-Experts with Expert
  Choice (https://arxiv.org/abs/2202.09368): each expert selects its top
  expert_capacity tokens. An individual token may be processed by multiple
  experts or none at all.
  """

  def _compute_routing_instructions(self, router_probs,
                                    expert_capacity):
    """Computes masks for the highest probability token per expert.

    Args:
      router_probs: <float32>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS]
        probabilities used to determine the routing of tokens to the experts.
      expert_capacity: Each group will send this many tokens to each expert.

    Returns:
      Dispatch and combine arrays for routing with masked matmuls.
    """
    tokens_per_group = router_probs.shape[1]

    # vmap over group dimension.
    router_probs_t = jax.vmap(lambda m: m.transpose())(router_probs)

    # Top expert_capacity router probability and corresponding token indices for
    # each expert. Shapes: [NUM_GROUPS, NUM_EXPERTS, EXPERT_CAPACITY].
    expert_gate, expert_index = _top_k(router_probs_t, k=expert_capacity)

    # Convert to one-hot mask of expert indices for each token in each group.
    # Shape: [NUM_GROUPS, NUM_EXPERTS, EXPERT_CAPACITY, TOKENS_PER_GROUP].
    dispatch_mask = jax.nn.one_hot(
        expert_index, tokens_per_group, dtype=jnp.int32)

    # Move axes to conform with shape expected by MoeLayer API.
    # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS, EXPERT_CAPACITY]
    dispatch_mask = jnp.moveaxis(dispatch_mask, 3, 1)

    # The combine array will be used for combining expert outputs, scaled by the
    # router probabilities. Shape: [NUM_GROUPS, NUM_EXPERTS, TOKENS_PER_GROUP,
    # EXPERT_CAPACITY].
    combine_array = jnp.einsum(
        "...ec,...tec->...tec",
        expert_gate,
        dispatch_mask,
        precision=jax.lax.Precision.DEFAULT)

    # Return to default dtype now that router computation is complete.
    combine_array = jax.lax.convert_element_type(combine_array, self.dtype)

    # Each expert is choosing tokens until it reaches full capacity, so we don't
    # need an auxiliary loading balancing loss for expert choice routing.
    auxiliary_loss = 0.0

    return RouterMask(dispatch_mask, combine_array, auxiliary_loss)

def _load_balancing_loss(router_probs,
                         expert_indices):
  """Computes auxiliary load balancing loss as in Switch Transformer.

  See Switch Transformer (https://arxiv.org/abs/2101.03961). This function
  implements the loss function presented in equations (4) - (6). It aims to
  penalize those cases where the routing between experts is unbalanced.

  Args:
    router_probs: Probability assigned to each expert per token. Shape:
      <float32>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS].
    expert_indices: <int>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS]
      indices identifying the top NUM_SELECTED_EXPERTS for a given token.

  Returns:
    The auxiliary loss.
  """
  num_experts = router_probs.shape[-1]

  # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_SELECTED_EXPERTS, NUM_EXPERTS].
  expert_mask = jax.nn.one_hot(expert_indices, num_experts, dtype=jnp.int32)
  # For a given token, determine if it was routed to a given expert.
  # Shape: [NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS]
  expert_mask = jnp.max(expert_mask, axis=-2)

  tokens_per_group_and_expert = jnp.mean(
      expert_mask, dtype=jnp.float32, axis=-2)
  router_prob_per_group_and_expert = jnp.mean(
      router_probs, dtype=jnp.float32, axis=-2)
  return jnp.mean(
      tokens_per_group_and_expert * router_prob_per_group_and_expert,
      dtype=jnp.float32) * num_experts**2

def _router_z_loss(router_logits):
  """Compute router z-loss.

   The router z-loss was introduced in Designing Effective Sparse Expert Models
   (https://arxiv.org/abs/2202.08906). It encourages router logits to remain
   small in an effort to improve stability.

  Args:
    router_logits: <float>[NUM_GROUPS, TOKENS_PER_GROUP, NUM_EXPERTS] router
      logits.

  Returns:
    Scalar router z-loss.
  """
  num_groups, tokens_per_group, _ = router_logits.shape
  log_z = jax.nn.logsumexp(router_logits, axis=-1)
  z_loss = log_z**2
  return jnp.sum(z_loss, dtype=jnp.float32) / (num_groups * tokens_per_group)


import numpy as np
import time

import torch
import torch.nn as nn
from torch.autograd import Function
from torch.cuda.amp import custom_bwd, custom_fwd

try:
    import _raymarching as _backend
except ImportError:
    from .backend import _backend

# ----------------------------------------
# utils
# ----------------------------------------

class _near_far_from_aabb(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, rays_o, rays_d, aabb, min_near=0.2):
        ''' near_far_from_aabb, CUDA implementation
        Calculate rays' intersection time (near and far) with aabb
        Args:
            rays_o: float, [N, 3]
            rays_d: float, [N, 3]
            aabb: float, [6], (xmin, ymin, zmin, xmax, ymax, zmax)
            min_near: float, scalar
        Returns:
            nears: float, [N]
            fars: float, [N]
        '''
        if not rays_o.is_cuda: rays_o = rays_o.cuda()
        if not rays_d.is_cuda: rays_d = rays_d.cuda()

        rays_o = rays_o.contiguous().view(-1, 3)
        rays_d = rays_d.contiguous().view(-1, 3)

        N = rays_o.shape[0] # num rays

        nears = torch.empty(N, dtype=rays_o.dtype, device=rays_o.device)
        fars = torch.empty(N, dtype=rays_o.dtype, device=rays_o.device)

        _backend.near_far_from_aabb(rays_o, rays_d, aabb, N, min_near, nears, fars)

        return nears, fars

near_far_from_aabb = _near_far_from_aabb.apply

# ----------------------------------------
# train functions
# ----------------------------------------

class _march_rays_train(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, rays_o, rays_d, bound, density_grid, mean_density, nears, fars, step_counter=None, mean_count=-1, perturb=False, align=-1, force_all_rays=False, dt_gamma=0, max_steps=1024):
        ''' march rays to generate points (forward only)
        Args:
            rays_o/d: float, [N, 3]
            bound: float, scalar
            density_grid: float, [C, H, H, H]
            mean_density: float, scalar
            nears/fars: float, [N]
            step_counter: int32, (2), used to count the actual number of generated points.
            mean_count: int32, estimated mean steps to accelerate training. (but will randomly drop rays if the actual point count exceeded this threshold.)
            perturb: bool
            align: int, pad output so its size is dividable by align, set to -1 to disable.
            force_all_rays: bool, ignore step_counter and mean_count, always calculate all rays. Useful if rendering the whole image, instead of some rays.
            dt_gamma: float, called cone_angle in instant-ngp, exponentially accelerate ray marching if > 0. (very significant effect, but generally lead to worse performance)
            max_steps: int, max number of sampled points along each ray, also affect min_stepsize.
        Returns:
            xyzs: float, [M, 3], all generated points' coords. (all rays concated, need to use `rays` to extract points belonging to each ray)
            dirs: float, [M, 3], all generated points' view dirs.
            deltas: float, [M, 2], all generated points' deltas. (first for RGB, second for Depth)
            rays: int32, [N, 3], all rays' (index, point_offset, point_count), e.g., xyzs[rays[i, 1]:rays[i, 2]] --> points belonging to rays[i, 0]
        '''

        if not rays_o.is_cuda: rays_o = rays_o.cuda()
        if not rays_d.is_cuda: rays_d = rays_d.cuda()
        if not density_grid.is_cuda: density_grid = density_grid.cuda()

        rays_o = rays_o.contiguous().view(-1, 3)
        rays_d = rays_d.contiguous().view(-1, 3)
        density_grid = density_grid.contiguous()

        N = rays_o.shape[0] # num rays

        C = density_grid.shape[0] # grid cascade
        H = density_grid.shape[1] # grid resolution

        M = N * max_steps # init max points number in total

        # running average based on previous epoch (mimic `measured_batch_size_before_compaction` in instant-ngp)
        # It estimate the max points number to enable faster training, but will lead to random ignored rays if underestimated.
        if not force_all_rays and mean_count > 0:
            if align > 0:
                mean_count += align - mean_count % align
            M = mean_count

        xyzs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
        dirs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
        deltas = torch.zeros(M, 2, dtype=rays_o.dtype, device=rays_o.device)
        rays = torch.empty(N, 3, dtype=torch.int32, device=rays_o.device) # id, offset, num_steps

        if step_counter is None:
            step_counter = torch.zeros(2, dtype=torch.int32, device=rays_o.device) # point counter, ray counter

        _backend.march_rays_train(rays_o, rays_d, density_grid, mean_density, bound, dt_gamma, max_steps, N, C, H, M, nears, fars, xyzs, dirs, deltas, rays, step_counter, perturb) # m is the actually used points number

        #print(step_counter, M)

        # only used at the first (few) epochs.
        if force_all_rays or mean_count <= 0:
            m = step_counter[0].item() # D2H copy
            if align > 0:
                m += align - m % align
            xyzs = xyzs[:m]
            dirs = dirs[:m]
            deltas = deltas[:m]

            torch.cuda.empty_cache()

        return xyzs, dirs, deltas, rays

march_rays_train = _march_rays_train.apply

class _composite_rays_train(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, sigmas, rgbs, deltas, rays):
        ''' composite rays' rgbs, according to the ray marching formula.
        Args:
            rgbs: float, [M, 3]
            sigmas: float, [M,]
            deltas: float, [M, 2]
            rays: int32, [N, 3]
        Returns:
            weights_sum: float, [N,], the alpha channel
            depth: float, [N, ], the Depth
            image: float, [N, 3], the RGB channel (after multiplying alpha!)
        '''

        sigmas = sigmas.contiguous()
        rgbs = rgbs.contiguous()

        M = sigmas.shape[0]
        N = rays.shape[0]

        weights_sum = torch.empty(N, dtype=sigmas.dtype, device=sigmas.device)
        depth = torch.empty(N, dtype=sigmas.dtype, device=sigmas.device)
        image = torch.empty(N, 3, dtype=sigmas.dtype, device=sigmas.device)

        _backend.composite_rays_train_forward(sigmas, rgbs, deltas, rays, M, N, weights_sum, depth, image)

        ctx.save_for_backward(sigmas, rgbs, deltas, rays, weights_sum, depth, image)
        ctx.dims = [M, N]

        return weights_sum, depth, image

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_weights_sum, grad_depth, grad_image):

        # NOTE: grad_depth is not used now! It won't be propagated to sigmas.

        grad_weights_sum = grad_weights_sum.contiguous()
        grad_image = grad_image.contiguous()

        sigmas, rgbs, deltas, rays, weights_sum, depth, image = ctx.saved_tensors
        M, N = ctx.dims

        grad_sigmas = torch.zeros_like(sigmas)
        grad_rgbs = torch.zeros_like(rgbs)

        _backend.composite_rays_train_backward(grad_weights_sum, grad_image, sigmas, rgbs, deltas, rays, weights_sum, image, M, N, grad_sigmas, grad_rgbs)

        return grad_sigmas, grad_rgbs, None, None

composite_rays_train = _composite_rays_train.apply

# ----------------------------------------
# infer functions
# ----------------------------------------

class _march_rays(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, n_alive, n_step, rays_alive, rays_t, rays_o, rays_d, bound, density_grid, mean_density, near, far, align=-1, perturb=False, dt_gamma=0, max_steps=1024):
        ''' march rays to generate points (forward only, for inference)
        Args:
            n_alive: int, number of alive rays
            n_step: int, how many steps we march
            rays_alive: int, [N], the alive rays' IDs in N (N >= n_alive, but we only use first n_alive)
            rays_t: float, [N], the alive rays' time, we only use the first n_alive.
            rays_o/d: float, [N, 3]
            bound: float, scalar
            density_grid: float, [C, H, H, H]
            mean_density: float, scalar
            nears/fars: float, [N]
            align: int, pad output so its size is dividable by align, set to -1 to disable.
            perturb: bool/int, int > 0 is used as the random seed.
            dt_gamma: float, called cone_angle in instant-ngp, exponentially accelerate ray marching if > 0. (very significant effect, but generally lead to worse performance)
            max_steps: int, max number of sampled points along each ray, also affect min_stepsize.
        Returns:
            xyzs: float, [n_alive * n_step, 3], all generated points' coords
            dirs: float, [n_alive * n_step, 3], all generated points' view dirs.
            deltas: float, [n_alive * n_step, 2], all generated points' deltas (here we record two deltas, the first is for RGB, the second for depth).
        '''

        if not rays_o.is_cuda: rays_o = rays_o.cuda()
        if not rays_d.is_cuda: rays_d = rays_d.cuda()

        rays_o = rays_o.contiguous().view(-1, 3)
        rays_d = rays_d.contiguous().view(-1, 3)

        C = density_grid.shape[0] # grid cascade
        H = density_grid.shape[1] # grid resolution
        M = n_alive * n_step

        if align > 0:
            M += align - (M % align)

        xyzs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
        dirs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
        deltas = torch.zeros(M, 2, dtype=rays_o.dtype, device=rays_o.device) # 2 vals, one for rgb, one for depth

        _backend.march_rays(n_alive, n_step, rays_alive, rays_t, rays_o, rays_d, bound, dt_gamma, max_steps, C, H, density_grid, mean_density, near, far, xyzs, dirs, deltas, perturb)

        return xyzs, dirs, deltas

march_rays = _march_rays.apply

class _composite_rays(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32) # need to cast sigmas & rgbs to float
    def forward(ctx, n_alive, n_step, rays_alive, rays_t, sigmas, rgbs, deltas, weights_sum, depth, image):
        ''' composite rays' rgbs, according to the ray marching formula. (for inference)
        Args:
            n_alive: int, number of alive rays
            n_step: int, how many steps we march
            rays_alive: int, [N], the alive rays' IDs in N (N >= n_alive, but we only use first n_alive)
            rays_t: float, [N], the alive rays' time, we only use the first n_alive.
            sigmas: float, [n_alive * n_step,]
            rgbs: float, [n_alive * n_step, 3]
            deltas: float, [n_alive * n_step, 2], all generated points' deltas (here we record two deltas, the first is for RGB, the second for depth).
        In-place Outputs:
            weights_sum: float, [N,], the alpha channel
            depth: float, [N,], the depth value
            image: float, [N, 3], the RGB channel (after multiplying alpha!)
        '''
        _backend.composite_rays(n_alive, n_step, rays_alive, rays_t, sigmas, rgbs, deltas, weights_sum, depth, image)
        return tuple()

composite_rays = _composite_rays.apply

class _compact_rays(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, n_alive, rays_alive, rays_alive_old, rays_t, rays_t_old, alive_counter):
        ''' compact rays, remove dead rays and reallocate alive rays, to accelerate next ray marching.
        Args:
            n_alive: int, number of alive rays
            rays_alive_old: int, [N]
            rays_t_old: float, [N], dead rays are marked by rays_t < 0
            alive_counter: int, [1], used to count remained alive rays.
        In-place Outputs:
            rays_alive: int, [N]
            rays_t: float, [N]
        '''
        _backend.compact_rays(n_alive, rays_alive, rays_alive_old, rays_t, rays_t_old, alive_counter)
        return tuple()

compact_rays = _compact_rays.apply


import cv2
import numpy as np

# Create a black image.
img = np.zeros((800, 800, 3), np.uint8)

# Initialize the Kalman filter.
kalman = cv2.KalmanFilter(4, 2)
kalman.measurementMatrix = np.array(
    [[1, 0, 0, 0],
     [0, 1, 0, 0]], np.float32)
kalman.transitionMatrix = np.array(
    [[1, 0, 1, 0],
     [0, 1, 0, 1],
     [0, 0, 1, 0],
     [0, 0, 0, 1]], np.float32)
kalman.processNoiseCov = np.array(
    [[1, 0, 0, 0],
     [0, 1, 0, 0],
     [0, 0, 1, 0],
     [0, 0, 0, 1]], np.float32) * 0.03

last_measurement = None
last_prediction = None

def on_mouse_moved(event, x, y, flags, param):
    global img, kalman, last_measurement, last_prediction

    measurement = np.array([[x], [y]], np.float32)
    if last_measurement is None:
        # This is the first measurement.
        # Update the Kalman filter's state to match the measurement.
        kalman.statePre = np.array(
            [[x], [y], [0], [0]], np.float32)
        kalman.statePost = np.array(
            [[x], [y], [0], [0]], np.float32)
        prediction = measurement
    else:
        kalman.correct(measurement)
        prediction = kalman.predict()  # Gets a reference, not a copy

        # Trace the path of the measurement in green.
        cv2.line(img, (last_measurement[0], last_measurement[1]),
                 (measurement[0], measurement[1]), (0, 255, 0))

        # Trace the path of the prediction in red.
        cv2.line(img, (last_prediction[0], last_prediction[1]),
                 (prediction[0], prediction[1]), (0, 0, 255))

    last_prediction = prediction.copy()
    last_measurement = measurement

cv2.namedWindow('kalman_tracker')
cv2.setMouseCallback('kalman_tracker', on_mouse_moved)

while True:
    cv2.imshow('kalman_tracker', img)
    k = cv2.waitKey(1)
    if k == 27:  # Escape
        cv2.imwrite('kalman.png', img)
        break


#!/usr/bin/env python

"""
This module provides classes to create phase diagrams.
"""

from __future__ import division

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "2.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__status__ = "Production"
__date__ = "Nov 25, 2012"

import collections
import numpy as np

from pyhull.convex_hull import ConvexHull

from pymatgen.core.composition import Composition
from pymatgen.phasediagram.entries import GrandPotPDEntry, TransformedPDEntry

from pymatgen.core.periodic_table import DummySpecie
from pymatgen.analysis.reaction_calculator import Reaction, ReactionError

class PhaseDiagram (object):
    """
    Simple phase diagram class taking in elements and entries as inputs.
    The algorithm is based on the work in the following papers:

    1. S. P. Ong, L. Wang, B. Kang, and G. Ceder, Li-Fe-P-O2 Phase Diagram from
       First Principles Calculations. Chem. Mater., 2008, 20(5), 1798-1807.
       doi:10.1021/cm702327g

    2. S. P. Ong, A. Jain, G. Hautier, B. Kang, G. Ceder, Thermal stabilities
       of delithiated olivine MPO4 (M=Fe, Mn) cathodes investigated using first
       principles calculations. Electrochem. Comm., 2010, 12(3), 427-430.
       doi:10.1016/j.elecom.2010.01.010

    .. attribute: elements:

        Elements in the phase diagram.

    ..attribute: all_entries

        All entries provided for Phase Diagram construction. Note that this
        does not mean that all these entries are actually used in the phase
        diagram. For example, this includes the positive formation energy
        entries that are filtered out before Phase Diagram construction.

    .. attribute: qhull_data

        Data used in the convex hull operation. This is essentially a matrix of
        composition data and energy per atom values created from qhull_entries.

    .. attribute: dim

        The dimensionality of the phase diagram.

    .. attribute: facets

        Facets of the phase diagram in the form of  [[1,2,3],[4,5,6]...]

    .. attribute: el_refs:

        List of elemental references for the phase diagrams. These are
        entries corresponding to the lowest energy element entries for simple
        compositional phase diagrams.

    .. attribute: qhull_entries:

        Actual entries used in convex hull. Excludes all positive formation
        energy entries.
    """

    # Tolerance for determining if formation energy is positive.
    formation_energy_tol = 1e-11

    def __init__(self, entries, elements=None):
        """
        Standard constructor for phase diagram.

        Args:
            entries:
                A list of PDEntry-like objects having an energy,
                energy_per_atom and composition.
            elements:
                Optional list of elements in the phase diagram. If set to None,
                the elements are determined from the the entries themselves.
        """
        if elements is None:
            elements = set()
            map(elements.update, [entry.composition.elements
                                  for entry in entries])
        elements = list(elements)
        # Qhull seems to be sensitive to choice of independent composition
        # components due to numerical issues in higher dimensions. The
        # code permutes the element sequence until one that works is found.
        dim = len(elements)
        el_refs = {}
        for el in elements:
            el_entries = filter(lambda e: e.composition.is_element and
                                e.composition.elements[0] == el, entries)
            if len(el_entries) == 0:
                raise PhaseDiagramError(
                    "There are no entries associated with terminal {}."
                    .format(el))
            el_refs[el] = min(el_entries, key=lambda e: e.energy_per_atom)

        data = []
        for entry in entries:
            comp = entry.composition
            row = map(comp.get_atomic_fraction, elements)
            row.append(entry.energy_per_atom)
            data.append(row)

        data = np.array(data)
        self.all_entries_hulldata = data[:, 1:]

        # Calculate formation energies and remove positive formation energy
        # entries
        vec = [el_refs[el].energy_per_atom for el in elements] + [-1]
        form_e = -np.dot(data, vec)
        ind = np.where(form_e <= -self.formation_energy_tol)[0].tolist()
        ind.extend(map(entries.index, el_refs.values()))
        qhull_entries = [entries[i] for i in ind]
        qhull_data = data[ind][:, 1:]

        if len(qhull_data) == dim:
            self.facets = [range(dim)]
        else:
            facets = ConvexHull(qhull_data, joggle=True).vertices
            finalfacets = []
            for facet in facets:
                is_non_element_facet = any(
                    (len(qhull_entries[i].composition) > 1 for i in facet))
                if is_non_element_facet:
                    m = qhull_data[facet]
                    m[:, -1] = 1
                    if abs(np.linalg.det(m)) > 1e-8:
                        finalfacets.append(facet)
            self.facets = finalfacets

        self.all_entries = entries
        self.qhull_data = qhull_data
        self.dim = dim
        self.el_refs = el_refs
        self.elements = elements
        self.qhull_entries = qhull_entries

    @property
    def unstable_entries(self):
        """
        Entries that are unstable in the phase diagram. Includes positive
        formation energy entries.
        """
        return [e for e in self.all_entries if e not in self.stable_entries]

    @property
    def stable_entries(self):
        """
        Returns the stable entries in the phase diagram.
        """
        stable_entries = set()
        for facet in self.facets:
            for vertex in facet:
                stable_entries.add(self.qhull_entries[vertex])
        return stable_entries

    def get_form_energy(self, entry):
        """
        Returns the formation energy for an entry (NOT normalized) from the
        elemental references.

        Args:
            entry:
                A PDEntry-like object.

        Returns:
            Formation energy from the elemental references.
        """
        comp = entry.composition
        energy = entry.energy - sum([comp[el] *
                                     self.el_refs[el].energy_per_atom
                                     for el in comp.elements])
        return energy

    def get_form_energy_per_atom(self, entry):
        """
        Returns the formation energy per atom for an entry from the
        elemental references.

        Args:
            entry:
                An PDEntry-like object

        Returns:
            Formation energy **per atom** from the elemental references.
        """
        comp = entry.composition
        return self.get_form_energy(entry) / comp.num_atoms

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        symbols = [el.symbol for el in self.elements]
        output = ["{} phase diagram".format("-".join(symbols)),
                  "{} stable phases: ".format(len(self.stable_entries)),
                  ", ".join([entry.name
                             for entry in self.stable_entries])]
        return "\n".join(output)

class GrandPotentialPhaseDiagram(PhaseDiagram):
    """
    A class representing a Grand potential phase diagram. Grand potential phase
    diagrams are essentially phase diagrams that are open to one or more
    components. To construct such phase diagrams, the relevant free energy is
    the grand potential, which can be written as the Legendre transform of the
    Gibbs free energy as follows

    Grand potential = G - u\ :sub:`X` N\ :sub:`X`\

    The algorithm is based on the work in the following papers:

    1. S. P. Ong, L. Wang, B. Kang, and G. Ceder, Li-Fe-P-O2 Phase Diagram from
       First Principles Calculations. Chem. Mater., 2008, 20(5), 1798-1807.
       doi:10.1021/cm702327g

    2. S. P. Ong, A. Jain, G. Hautier, B. Kang, G. Ceder, Thermal stabilities
       of delithiated olivine MPO4 (M=Fe, Mn) cathodes investigated using first
       principles calculations. Electrochem. Comm., 2010, 12(3), 427-430.
       doi:10.1016/j.elecom.2010.01.010
    """

    def __init__(self, entries, chempots, elements=None):
        """
        Standard constructor for grand potential phase diagram.

        Args:
            entries:
                A list of PDEntry-like objects having an energy,
                energy_per_atom and composition.
            chempots:
                A dict of {element: float} to specify the chemical potentials
                of the open elements.
            elements:
                Optional list of elements in the phase diagram. If set to None,
                the elements are determined from the entries themselves.
        """
        if elements is None:
            elements = set()
            map(elements.update, [entry.composition.elements
                                  for entry in entries])

        elements = set(elements).difference(chempots.keys())
        all_entries = [GrandPotPDEntry(e, chempots)
                       for e in entries
                       if (not e.is_element) or
                       e.composition.elements[0] in elements]
        self.chempots = chempots

        super(GrandPotentialPhaseDiagram, self).__init__(all_entries, elements)

    def __str__(self):
        output = []
        chemsys = "-".join([el.symbol for el in self.elements])
        output.append("{} grand potential phase diagram with ".format(chemsys))
        output[-1] += ", ".join(["u{}={}".format(el, v)
                                 for el, v in self.chempots.items()])
        output.append("{} stable phases: ".format(len(self.stable_entries)))
        output.append(", ".join([entry.name
                                 for entry in self.stable_entries]))
        return "\n".join(output)

class CompoundPhaseDiagram(PhaseDiagram):
    """
    Generates phase diagrams from compounds as terminations instead of
    elements.
    """

    # Tolerance for determining if amount of a composition is positive.
    amount_tol = 1e-5

    def __init__(self, entries, terminal_compositions,
                 normalize_terminal_compositions=True):
        """
        Args:
            entries:
                Sequence of input entries. For example, if you want a Li2O-P2O5
                phase diagram, you might have all Li-P-O entries as an input.
            terminal_compositions:
                Terminal compositions of phase space. In the Li2O-P2O5 example,
                these will be the Li2O and P2O5 compositions.
            normalize_terminal_compositions:
                Whether to normalize the terminal compositions to a per atom
                basis. If normalized, the energy above hulls will be consistent
                for comparison across systems. Non-normalized terminals are
                more intuitive in terms of compositional breakdowns.
        """
        self.original_entries = entries
        self.terminal_compositions = terminal_compositions
        self.normalize_terminals = normalize_terminal_compositions
        (pentries, species_mapping) = \
            self.transform_entries(entries, terminal_compositions)
        self.species_mapping = species_mapping
        PhaseDiagram.__init__(self, pentries,
                              elements=species_mapping.values())

    def transform_entries(self, entries, terminal_compositions):
        """
        Method to transform all entries to the composition coordinate in the
        terminal compositions. If the entry does not fall within the space
        defined by the terminal compositions, they are excluded. For example,
        Li3PO4 is mapped into a Li2O:1.5, P2O5:0.5 composition. The terminal
        compositions are represented by DummySpecies.

        Args:
            entries:
                Sequence of all input entries
            terminal_compositions:
                Terminal compositions of phase space.

        Returns:
            Sequence of TransformedPDEntries falling within the phase space.
        """
        new_entries = []
        if self.normalize_terminals:
            fractional_comp = [c.get_fractional_composition()
                               for c in terminal_compositions]
        else:
            fractional_comp = terminal_compositions

        #Map terminal compositions to unique dummy species.
        sp_mapping = collections.OrderedDict()
        for i, comp in enumerate(fractional_comp):
            sp_mapping[comp] = DummySpecie("X" + chr(102 + i))

        for entry in entries:
            try:
                rxn = Reaction(fractional_comp, [entry.composition])
                rxn.normalize_to(entry.composition)
                #We only allow reactions that have positive amounts of
                #reactants.
                if all([rxn.get_coeff(comp) <= CompoundPhaseDiagram.amount_tol
                        for comp in fractional_comp]):
                    newcomp = {sp_mapping[comp]: -rxn.get_coeff(comp)
                               for comp in fractional_comp}
                    newcomp = {k: v for k, v in newcomp.items()
                               if v > CompoundPhaseDiagram.amount_tol}
                    transformed_entry = \
                        TransformedPDEntry(Composition(newcomp), entry)
                    new_entries.append(transformed_entry)
            except ReactionError:
                #If the reaction can't be balanced, the entry does not fall
                #into the phase space. We ignore them.
                pass
        return new_entries, sp_mapping

class PhaseDiagramError(Exception):
    """
    An exception class for Phase Diagram generation.
    """
    pass


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Class definitions for RRTMG objects in RCE model. Implements bridge between
pyRRTMG API and the RCE model solver requirements.

Created on Tue Nov 15 18:43:23 2016

@author: maxwell
"""

from radiation.radiation import RadModel,Flux
import radiation.pyrrtmg as rr
import numpy as np

class RRTMGModel(RadModel):
    """
    Class construct for pyrrtmg
    """

    def __init__(self,cpdair=None,**kwargs):

        print('ititializing rrtmg object')
        if cpdair is None:
            print(
                "WARNING: cpdair not provided to RRTMGModel. " ,
                "Using pyRRTMG default."
                 )
            rr.lw.init()
            rr.sw.init()
        else:
            rr.lw.init(cpdair)
            rr.sw.init(cpdair)
    def radiation(self,atms,cparm,lwparm, swparm):
        """
        Calculates RRTMG radiation.

        Expects an Atmosphere object with the following variables (size):
            play (n)
            plev (n+1)
            tlay (n)
            tlev (n+1)
            qlay (n)
            o3lay (n)
        Expects a ChemParm object with the following scalars:
            co2
            ch4
            n2o
            o2
            cfc11
            cfc12
            cfc22
            ccl4
        Expects a LWParm object with the following scalars:
            emis
        Expects a SWParm object with the following scalars:
            albedo
            fday
            coszen
            scon

        Returns a Flux object with short and longwave fluxes
        """
        fuir = np.empty(len(atms.plev))
        fdir = np.empty(len(atms.plev))
        fusw = np.empty(len(atms.plev))
        fdsw = np.empty(len(atms.plev))

        # flip indices, since rrtm expects pressure to decrease with index,
        # while the Atmosphere class expects pressure to increase with index.
        (fuir[::-1], fdir[::-1]) = rr.lw.rad(atms.play[::-1], atms.plev[::-1],
                                 atms.tlay[::-1], atms.tlev[::-1],
                                 atms.tsfc,
                                 atms.qlay[::-1], atms.o3lay[::-1],
                                 **cparm, **lwparm)
        (fusw[::-1], fdsw[::-1]) = rr.sw.rad(atms.play[::-1], atms.plev[::-1],
                                 atms.tlay[::-1], atms.tlev[::-1],
                                 atms.tsfc,
                                 atms.qlay[::-1], atms.o3lay[::-1],
                                 **cparm, **swparm)

        return Flux(fuir,fdir,fusw,fdsw)