import random
from typing import Any

import numpy as np
from process_bigraph import Process
from simularium_readdy_models.actin import (
    ActinSimulation,
    ActinGenerator,
    FiberData,
)
from simularium_readdy_models.common import ReaddyUtil, get_membrane_monomers


def get_monomers():
    actin_monomers = ActinGenerator.get_monomers(
        fibers_data=[
            FiberData(
                28,
                [
                    np.array([-25, 0, 0]),
                    np.array([25, 0, 0]),
                ],
                "Actin-Polymer",
            )
        ],
        use_uuids=False,
        start_normal=np.array([0.0, 1.0, 0.0]),
        longitudinal_bonds=True,
        barbed_binding_site=True,
    )
    actin_monomers = ActinGenerator.setup_fixed_monomers(
        actin_monomers,
        orthogonal_seed=True,
        n_fixed_monomers_pointed=3,
        n_fixed_monomers_barbed=0,
    )
    membrane_monomers = get_membrane_monomers(
        center=np.array([25.0, 0.0, 0.0]),
        size=np.array([0.0, 100.0, 100.0]),
        particle_radius=2.5,
        start_particle_id=len(actin_monomers["particles"].keys()),
        top_id=1,
    )
    free_actin_monomers = ActinGenerator.get_free_actin_monomers(
        concentration=500.0,
        box_center=np.array([12.0, 0.0, 0.0]),
        box_size=np.array([20.0, 50.0, 50.0]),
        start_particle_id=len(actin_monomers["particles"].keys())
        + len(membrane_monomers["particles"].keys()),
        start_top_id=2,
    )
    monomers = {
        "particles": {**actin_monomers["particles"], **membrane_monomers["particles"]},
        "topologies": {
            **actin_monomers["topologies"],
            **membrane_monomers["topologies"],
        },
    }
    monomers = {
        "particles": {**monomers["particles"], **free_actin_monomers["particles"]},
        "topologies": {**monomers["topologies"], **free_actin_monomers["topologies"]},
    }
    return monomers


class ReaddyActinMembrane(Process):
    """
    This process runs ReaDDy models with coarse-grained particle
    actin filaments and membrane patches.
    """

    config_schema = {
        "name": "string",
        "internal_timestep": "float",
        "box_size": "tuple[float,float,float]",
        "periodic_boundary": "boolean",
        "reaction_distance": "float",
        "n_cpu": "integer",
        "only_linear_actin_constraints": "boolean",
        "reactions": "boolean",
        "dimerize_rate": "float",
        "dimerize_reverse_rate": "float",
        "trimerize_rate": "float",
        "trimerize_reverse_rate": "float",
        "pointed_growth_ATP_rate": "float",
        "pointed_growth_ADP_rate": "float",
        "pointed_shrink_ATP_rate": "float",
        "pointed_shrink_ADP_rate": "float",
        "barbed_growth_ATP_rate": "float",
        "barbed_growth_ADP_rate": "float",
        "nucleate_ATP_rate": "float",
        "nucleate_ADP_rate": "float",
        "barbed_shrink_ATP_rate": "float",
        "barbed_shrink_ADP_rate": "float",
        "arp_bind_ATP_rate": "float",
        "arp_bind_ADP_rate": "float",
        "arp_unbind_ATP_rate": "float",
        "arp_unbind_ADP_rate": "float",
        "barbed_growth_branch_ATP_rate": "float",
        "barbed_growth_branch_ADP_rate": "float",
        "debranching_ATP_rate": "float",
        "debranching_ADP_rate": "float",
        "cap_bind_rate": "float",
        "cap_unbind_rate": "float",
        "hydrolysis_actin_rate": "float",
        "hydrolysis_arp_rate": "float",
        "nucleotide_exchange_actin_rate": "float",
        "nucleotide_exchange_arp_rate": "float",
        "verbose": "boolean",
        "use_box_actin": "boolean",
        "use_box_arp": "boolean",
        "use_box_cap": "boolean",
        "obstacle_radius": "float",
        "obstacle_diff_coeff": "float",
        "use_box_obstacle": "boolean",
        "position_obstacle_stride": "integer",
        "displace_pointed_end_tangent": "boolean",
        "displace_pointed_end_radial": "boolean",
        "tangent_displacement_nm": "float",
        "radial_displacement_radius_nm": "float",
        "radial_displacement_angle_deg": "float",
        "longitudinal_bonds": "boolean",
        "displace_stride": "integer",
        "bonds_force_multiplier": "float",
        "angles_force_constant": "float",
        "dihedrals_force_constant": "float",
        "actin_constraints": "boolean",
        "use_box_actin": "boolean",
        "actin_box_center_x": "float",
        "actin_box_center_y": "float",
        "actin_box_center_z": "float",
        "actin_box_size_x": "float",
        "actin_box_size_y": "float",
        "actin_box_size_z": "float",
        "add_extra_box": "boolean",
        "barbed_binding_site": "boolean",
        "binding_site_reaction_distance": "float",
        "add_membrane": "boolean",
        "membrane_center_x": "float",
        "membrane_center_y": "float",
        "membrane_center_z": "float",
        "membrane_size_x": "float",
        "membrane_size_y": "float",
        "membrane_size_z": "float",
        "membrane_particle_radius": "float",
        "obstacle_controlled_position_x": "float",
        "obstacle_controlled_position_y": "float",
        "obstacle_controlled_position_z": "float",
        "random_seed": "integer",
        "total_steps": "float",
        "actin_concentration": "integer",  # 0
        "arp23_concentration": "integer",  # 0
        "cap_concentration": "integer",  # 0
        "n_fixed_monomers_barbed": "integer",  # 0
        "tangent_displace_speed_um_s": "float",
        "plot_actin_compression": "boolean",  # true
        "visualize_edges": "boolean",  # true
        "visualize_normals": "boolean",  # true
        "visualize_control_pts": "boolean",  # true
    }

    def initialize(self, config, readdy_system=None):
        random.seed(self.config["random_seed"])
        np.random.seed(self.config["random_seed"])
        actin_simulation = ActinSimulation(self.config, False, False, readdy_system)
        self.readdy_system = actin_simulation.system
        self.readdy_simulation = actin_simulation.simulation

    def initial_state(self):
        return get_monomers()

    def inputs(self):
        return {
            "topologies": "map[topology]",
            "particles": "map[particle]",
        }

    def outputs(self):
        return {
            "topologies": "map[topology]",
            "particles": "map[particle]",
        }

    def update(self, inputs, interval):
        self.initialize(self.config, self.readdy_system)

        ReaddyUtil.add_monomers_from_data(self.readdy_simulation, inputs)

        simulate_readdy(
            self.config["internal_timestep"],
            self.readdy_system,
            self.readdy_simulation,
            interval,
        )

        id_diff = id_difference(self.readdy_simulation.current_topologies)

        readdy_monomers = ReaddyUtil.get_current_monomers(
            self.readdy_simulation.current_topologies, id_diff
        )

        return readdy_monomers


def id_difference(current_topologies):
    """
    Get the first ID coming out of ReaDDy, it should be zero
    unless Readdy ran multiple times and cached the IDs, which
    would cause Vivarium to create new particles instead of
    updating existing particles.

    (This is a HACK needed as long as ReaDDy has this behavior.)
    """
    return current_topologies[0].particles[0].id


def simulate_readdy(internal_timestep, readdy_system, readdy_simulation, timestep):
    """
    Simulate in ReaDDy for the given timestep
    """

    def loop():
        readdy_actions = readdy_simulation._actions
        init = readdy_actions.initialize_kernel()
        diffuse = readdy_actions.integrator_euler_brownian_dynamics(internal_timestep)
        calculate_forces = readdy_actions.calculate_forces()
        create_nl = readdy_actions.create_neighbor_list(
            readdy_system.calculate_max_cutoff().magnitude
        )
        update_nl = readdy_actions.update_neighbor_list()
        react = readdy_actions.reaction_handler_uncontrolled_approximation(
            internal_timestep
        )
        init()
        create_nl()
        calculate_forces()
        update_nl()
        n_steps = int(timestep / internal_timestep)
        print(f"running readdy for {n_steps} steps")
        for t in range(1, n_steps + 1):
            diffuse()
            update_nl()
            react()
            update_nl()
            calculate_forces()

    readdy_simulation._run_custom_loop(loop, show_summary=False)
