"""nodes.py: Node abstract class (or interface)"""

from __future__ import annotations
from copy import deepcopy
from abc import ABC, abstractmethod
from itertools import product

class Node(ABC):
    """
    Node class, an element of a tree
    """
    def __init__(self, function: callable, number_of_arguments: int):
        if not hasattr(function, "__call__"):
            TypeError("Function {} given is not a function".format(function))
        self.function = function
        self.number_of_arguments = number_of_arguments
        self.arguments = list()
        self._is_node = True
        self.type = None

    @abstractmethod
    def evaluate(self, **kwargs) -> object or [object]:
        """
        Evaluate recursive all nodes child of this one

        :return: Value
        """
        pass

    def replace_node(self, other: Node, position: int) -> None:
        """
        Replace one nodes with another

        :param other: Node to replace this one
        :param position: Position of nodes to replace
        :return: None (replace nodes itself)
        """
        try:
            self.arguments[position] = deepcopy(other)
        except IndexError:
            raise AttributeError("No nodes on the {} position, "
                                 "this nodes has {} sub nodes".format(position, self.number_of_arguments))

    def __eq__(self, other: Node) -> bool:
        """For test purpose"""
        try:
            if not other._is_node:
                return False
        except AttributeError:
            return False
        if self.function != other.function:
            return False
        if self.number_of_arguments != other.number_of_arguments:
            return False
        for arg1, arg2 in zip(self.arguments, other.arguments):
            if arg1 != arg2:
                return False
        return True

    def __repr__(self) -> str:
        children = ""
        for child in self.arguments:
            children += "\n|-{}".format(str(child))
        return "({})" + children.replace("\n", "\n  ")

    @classmethod
    @abstractmethod
    def get_arguments(cls) -> int:
        """
        Get number of arguments that has to have
        """
        pass

    @staticmethod
    def __prepare_values(values: {list}) -> ([str], [tuple]):
        """
        Prepares values for be used for variable terminal

        :param values: Values as dictionary of list, for facilitate use
        :return: Values process: list of str (variables provided)
                list of tuples (values of variables combined in order of variables)
        """
        ready_values = product(*values.values())
        return list(values.keys()), list(ready_values)

    @staticmethod
    def _receive_values(**kwargs) -> (bool, [str], [tuple]):
        """
        Receive values, process if not process yet

        :param kwargs:
        :return:
        """
        try:
            if kwargs["ready"]:
                variables, values = kwargs["values"]
                return True, variables, values
        except KeyError:
            try:
                variables, values = Node.__prepare_values(kwargs["values"])
                return True, variables, values
            except KeyError:
                raise ValueError("Not values expected to test")


import os
import inspect
import time
import matplotlib.pyplot as plt
import numpy as np
import pybullet as p

class joint(object):
    def __init__(self,Uid,joint_num):
        self.Uid = Uid
        self.joint_num = joint_num
        p.enableJointForceTorqueSensor(self.Uid, self.joint_num, 1)
        self.joint_force_fx = []
        self.joint_force_fy = []
        self.joint_force_fz = []
        self.joint_force_Mx = []
        self.joint_force_My = []
        self.joint_force_Mz = []
        self.joint_motor_torque = []
        self.joint_vel = []

    def append_state(self):
        self.state = p.getJointState(self.Uid, self.joint_num)
        self.joint_force_fx.append(self.state[2][0])
        self.joint_force_fy.append(self.state[2][1])
        self.joint_force_fz.append(self.state[2][2])
        self.joint_force_Mx.append(self.state[2][3])
        self.joint_force_My.append(self.state[2][4])
        self.joint_force_Mz.append(self.state[2][5])
        self.joint_motor_torque.append(self.state[3])
        self.joint_vel.append(self.state[1])

    def plot(self):
        self.joint_force_fx = np.array(self.joint_force_fx)
        self.joint_force_fy = np.array(self.joint_force_fy)
        self.joint_force_fz = np.array(self.joint_force_fz)
        self.joint_force_Mx = np.array(self.joint_force_Mx)
        self.joint_force_My = np.array(self.joint_force_My)
        self.joint_force_Mz = np.array(self.joint_force_Mz)
        self.joint_motor_torque = np.array(self.joint_motor_torque)
        self.joint_vel = np.array(self.joint_vel)
        self.steps = range(len(self.joint_force_fx))
        fig = plt.figure()
        fig.suptitle('figure of joint %s' % self.joint_num)
        ax1 = fig.add_subplot(2,2,1)
        ax1.set_title('force x,y,z')
        ax1.plot(self.steps, self.joint_force_fx, label='fx')
        ax1.plot(self.steps, self.joint_force_fy, label='fy')
        ax1.plot(self.steps, self.joint_force_fz, label='fz')
        ax1.legend()

        ax2 = fig.add_subplot(2,2,2)
        ax2.set_title('force Mx,My,Mz')
        ax2.plot(self.steps, self.joint_force_Mx, label='Mx')
        ax2.plot(self.steps, self.joint_force_My, label='My')
        ax2.plot(self.steps, self.joint_force_Mz, label='Mz')
        ax2.legend()

        ax3 = fig.add_subplot(2,2,3)
        ax3.set_title('Motor torque')
        ax3.plot(self.steps, self.joint_motor_torque,label='motor_torque')
        ax3.legend()

        ax4 = fig.add_subplot(2, 2, 4)
        ax4.set_title('joint_velocity')
        ax4.plot(self.steps, self.joint_vel, label='joint_velocity')
        ax4.legend()

currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(os.path.dirname(currentdir))
os.sys.path.insert(0, parentdir)
physicsClient = p.connect(p.GUI)
p.setGravity(0, 0, -9.8)
planeUid = p.loadURDF(os.path.join(parentdir, "My_env_test/neobotix_schunk_pybullet/data/plane.urdf"))
neobotixschunkUid = p.loadURDF(
    os.path.join(parentdir, "My_env_test/neobotix_schunk_pybullet/data/neobotixschunk/mp500lwa4d_test.urdf"),
    useFixedBase=False, flags=p.URDF_USE_SELF_COLLISION)

joint_observ1 = joint(neobotixschunkUid, 7)
joint_observ2 = joint(neobotixschunkUid, 9)
#
quitId = p.addUserDebugParameter("quit", 0, 1, 0)
#
p.setJointMotorControl2(neobotixschunkUid,9,p.POSITION_CONTROL,
                       targetPosition=0,force=50) #fz的方向为arm7旋转轴，因此arm6旋转时，会使重力在fz上的分力减小
# p.setJointMotorControl2(neobotixschunkUid,9,p.POSITION_CONTROL,
#                        targetPosition=0,force=10)   #力是表示在child_link坐标系上的，因此随着旋转fx，fy会变化
# p.setJointMotorControl2(neobotixschunkUid,9,p.POSITION_CONTROL,
#                        targetPosition=0,force=5)
# p.setJointMotorControl2(neobotixschunkUid,8,p.POSITION_CONTROL,
#                        targetPosition=0,force=5)
# p.setJointMotorControl2(neobotixschunkUid,7,p.POSITION_CONTROL,
#                        targetPosition=0,force=5)
# p.setJointMotorControl2(neobotixschunkUid,6,p.POSITION_CONTROL,
#                        targetPosition=0,force=5)
# p.setJointMotorControl2(neobotixschunkUid,5,p.POSITION_CONTROL,
#                        targetPosition=0,force=5)

num_joint = p.getNumJoints(neobotixschunkUid)
for i in range(num_joint):
    print('joint_info:', p.getJointInfo(neobotixschunkUid, i))  #joint_info return linkName of the Childlink of the joint
    print('Dynamik of links:', p.getDynamicsInfo(neobotixschunkUid, i))
    link_state = p.getLinkState(neobotixschunkUid, i)
    print('links position:', link_state[0])

pos_z = []

# while 1:
#     q = p.readUserDebugParameter(quitId)
#     if q>0:
#         break
#     p.stepSimulation()
#     joint_observ1.append_state()
#     time.sleep(1/240)

for i in range(200):
    q = p.readUserDebugParameter(quitId)
    if q>0:
        break
    p.stepSimulation()
    base_pos = p.getBasePositionAndOrientation(neobotixschunkUid)[0]
    pos_z.append(base_pos[2])
    joint_observ1.append_state()
    joint_observ2.append_state()
    time.sleep(1/240)

p.setJointMotorControl2(neobotixschunkUid,9,p.POSITION_CONTROL,
                       targetPosition=-0.5,force=50)
for i in range(200):
    q = p.readUserDebugParameter(quitId)
    if q>0:
        break
    p.stepSimulation()
    base_pos = p.getBasePositionAndOrientation(neobotixschunkUid)[0]
    pos_z.append(base_pos[2])
    joint_observ1.append_state()
    joint_observ2.append_state()
    time.sleep(1/240)

p.disconnect()

pos_z = np.array(pos_z)
step = np.arange(0, len(pos_z), 1)
plt.figure()
plt.suptitle('z position of base')
plt.plot(step, pos_z, 'r')
plt.legend()

joint_observ1.plot()
joint_observ2.plot()
plt.show()


import os
from textwrap import dedent

from nose.tools import eq_

from douglas.plugins import pages
from douglas.tests import PluginTest

class PagesTest(PluginTest):
    def setUp(self):
        PluginTest.setUp(self, pages)

    def test_is_frontpage(self):
        # test setup-related is_frontpage = False possibilities
        eq_(pages.is_frontpage({}, {}), False)
        eq_(pages.is_frontpage({'PATH_INFO': '/'}, {}), False)
        eq_(pages.is_frontpage({'PATH_INFO': '/'}, {'pages_frontpage': False}),
            False)

        # test path-related possibilities
        for path, expected in (('/', True),
                               ('/index', True),
                               ('/index.html', True),
                               ('/index.xml', True),
                               ('/foo', False)):
            eq_(pages.is_frontpage({'PATH_INFO': path},
                                   {'pages_frontpage': True}),
                expected)

    def test_pages(self):
        self.create_file('pages/about.txt', dedent("""\
        About this blog
        #meta1 val1
        body
        """))

        extra_cfg = {
            'pagesdir': os.path.join(self.datadir, 'pages')
        }
        req = self.build_request(extra_cfg)

        req.get_http()['PATH_INFO'] = '/pages/about'
        entries = pages.cb_filelist({'request': req})

        eq_(len(entries), 1)
        eq_(entries[0]['title'], 'About this blog')


"""
An analyzer that collects semantics at each program point.
"""

from collections import defaultdict
from xml.sax.saxutils import escape

from lalcheck.ai import domains
from lalcheck.ai.interpretations import def_provider_builder
from lalcheck.ai.irs.basic import visitors
from lalcheck.ai.irs.basic.purpose import SyntheticVariable
from lalcheck.ai.irs.basic.tree import Variable
from lalcheck.ai.utils import KeyCounter, concat_dicts
from lalcheck.tools import dot_printer
from lalcheck.tools.digraph import Digraph

from lalcheck.ai.irs.basic.tools import (
    CFGBuilder,
    ExprEvaluator,
    ExprSolver
)
from lalcheck.ai.irs.basic.tools import PrettyPrinter

def updated_state(state, var, value):
    i = var.data.index
    return state[:i] + (value,) + state[i+1:]

class _VarTracker(visitors.CFGNodeVisitor):
    def __init__(self, var_set, vars_domain, evaluator, c_solver):
        self.vars = var_set
        self.evaluator = evaluator
        self.constr_solver = c_solver
        self.vars_domain = vars_domain

    def visit_assign(self, assign, state):
        return updated_state(
            state,
            assign.id.var,
            self.evaluator.eval(assign.expr, state)
        )

    def visit_assume(self, assume, state):
        new_state, success = self.constr_solver.solve(assume.expr, state)
        return new_state if success else self.vars_domain.bottom

    def visit_read(self, read, state):
        return updated_state(
            state,
            read.id.var,
            self.evaluator.model[read.id.var].domain.top
        )

class _SimpleTraceLattice(domains.FiniteSubsetLattice):
    def __init__(self, *args):
        super(_SimpleTraceLattice, self).__init__(*args)

    def update(self, a, b, widen=False):
        return super(_SimpleTraceLattice, self).update(a, b, False)

class MergePredicateBuilder(object):
    def __init__(self, predicate):
        self.predicate = predicate

    def __or__(self, other):
        def f(trace_domain, vals_domain):
            s_pred = self.build(trace_domain, vals_domain)
            o_pred = other.build(trace_domain, vals_domain)

            return lambda a, b: s_pred(a, b) or o_pred(a, b)

        return MergePredicateBuilder(f)

    def __and__(self, other):
        def f(trace_domain, vals_domain):
            s_pred = self.build(trace_domain, vals_domain)
            o_pred = other.build(trace_domain, vals_domain)

            return lambda a, b: s_pred(a, b) and o_pred(a, b)

        return MergePredicateBuilder(f)

    def build(self, trace_domain, vals_domain):
        return self.predicate(trace_domain, vals_domain)

def _mp_always(*_):
    return lambda *_: True

def _mp_never(*_):
    return lambda *_: False

def _mp_le_traces(trace_domain, _):
    return lambda a, b: trace_domain.le(a[0], b[0])

def _mp_eq_vals(_, vals_domain):
    return lambda a, b: vals_domain.eq(a[1], b[1])

MergePredicateBuilder.Always = MergePredicateBuilder(_mp_always)
MergePredicateBuilder.Never = MergePredicateBuilder(_mp_never)
MergePredicateBuilder.Le_Traces = MergePredicateBuilder(_mp_le_traces)
MergePredicateBuilder.Eq_Vals = MergePredicateBuilder(_mp_eq_vals)

class ExternalCallStrategy(object):
    def __call__(self, sig):
        """
        Given a signature, returns an interpretation of this function as a
        pair which contains:
        - The forward implementation of that function.
        - The backward implementation of that function.

        :param lalcheck.ai.interpretations.Signature sig: The signature for
            which to provide an interpretation.

        :rtype: (*object->object, *object->object)
        """
        raise NotImplementedError

    def as_def_provider(self):
        return def_provider_builder(self.__call__)

class KnownTargetCallStrategy(ExternalCallStrategy):
    def __init__(self, progs):
        self.progs = progs

    def _get_provider(self, sig, prog):
        """
        Given a signature and the program that it designates, returns an
        interpretation of this function as a pair which contains:
        - The forward implementation of that function.
        - The backward implementation of that function.

        :param lalcheck.ai.interpretations.Signature sig: The signature for
            which to provide an interpretation.

        :param lalcheck.ai.irs.basic.tree.Program prog: The IR of the program
            that is designated by the signature.

        :rtype: (*object->object, *object->object)
        """
        raise NotImplementedError

    def __call__(self, sig):
        for prog in self.progs:
            if sig.name == prog.data.fun_id:
                return self._get_provider(sig, prog)

class UnknownTargetCallStrategy(ExternalCallStrategy):
    def __init__(self):
        super(UnknownTargetCallStrategy, self).__init__()

    def __call__(self, sig):
        def f(*_):
            if len(sig.out_param_indices) == 0:
                return (sig.output_domain.top
                        if sig.output_domain else ())

            return tuple(
                sig.input_domains[i].top
                for i in sig.out_param_indices
            ) + ((sig.output_domain.top,) if sig.output_domain else ())

        def inv(expected, *vals):
            if len(vals) == 1:
                return vals[0]
            else:
                return vals

        return f, inv

class TopDownCallStrategy(KnownTargetCallStrategy):
    def __init__(self, progs, get_model, get_merge_pred_builder):
        super(TopDownCallStrategy, self).__init__(progs)
        self.get_model = get_model
        self.get_merge_pred_builder = get_merge_pred_builder

    def _get_provider(self, sig, prog):
        def f(*args):
            arg_values = {
                param: value
                for param, value in zip(prog.data.param_vars, args)
            }

            prog_model = self.get_model(prog)

            analysis = compute_semantics(
                prog,
                prog_model,
                self.get_merge_pred_builder(),
                arg_values
            )

            # Get all environments that can result from the analysis of the
            # function called.
            envs = [
                values
                for leaf in analysis.cfg.leafs()
                for _, values in analysis.semantics[leaf].iteritems()
            ]

            # Fetch the indices of variables that are marked out.
            out_vars = [
                prog.data.param_vars[i]
                for i in sig.out_param_indices
            ]

            # Compute their value after the call using the environments.
            param_values = tuple(
                reduce(
                    prog_model[var].domain.join,
                    (env[var] for env in envs),
                    prog_model[var].domain.bottom
                )
                for var in out_vars
            )

            # Compute the value of the return variable (if any) after the call
            # using the environments.
            result_var = prog.data.result_var
            result_value = (reduce(
                prog_model[result_var].domain.join,
                (env[result_var] for env in envs),
                prog_model[result_var].domain.bottom
            ),) if result_var is not None else ()

            if len(out_vars) == 0 and result_var is not None:
                # We must return the value directly when there is a single
                # element to return instead of a 1-tuple.
                return result_value[0]
            else:
                return param_values + result_value

        def inv(*_):
            raise NotImplementedError

        return f, inv

def _html_render_node(node):
    return escape(PrettyPrinter.pretty_print(node))

def _save_cfg_to(file_name, cfg):
    def render_node(node):
        return (_html_render_node(node),) if node is not None else ()

    def render_widening_point(is_widening_point):
        return (escape('<widening_point>'),) if is_widening_point else ()

    with open(file_name, 'w') as f:
        f.write(dot_printer.gen_dot(cfg, [
            dot_printer.DataPrinter('node', render_node),
            dot_printer.DataPrinter('is_widening_point', render_widening_point)
        ]))

def _build_resulting_graph(file_name, cfg, results, trace_domain, model):
    paths = defaultdict(list)

    var_set = {
        v
        for node, state in results.iteritems()
        for trace, values in state.iteritems()
        for v, _ in values.iteritems()
    }

    for node, state in results.iteritems():
        for trace, values in state.iteritems():
            paths[frozenset(trace)].append(
                Digraph.Node(
                    node.name,
                    ___orig=node,
                    **{
                        v.name: value
                        for v, value in values.iteritems()
                        if not SyntheticVariable.is_purpose_of(v)
                    }
                )
            )

    edges = []
    for trace, nodes in paths.iteritems():
        for node in nodes:
            predecessors = [
                n
                for t, ns in paths.iteritems()
                for n in ns
                if trace_domain.le(t, trace)
                if n.data.___orig in cfg.ancestors(node.data.___orig)
            ]

            for pred in predecessors:
                edges.append(Digraph.Edge(pred, node))

    res_graph = Digraph(
        [n for _, nodes in paths.iteritems() for n in nodes],
        edges
    )

    def print_orig(orig):
        if orig.data.node is not None:
            return (
                '<i>{}</i>'.format(_html_render_node(orig.data.node)),
            )
        return ()

    def print_result_builder(v):
        return lambda value: (
            "{} ∈".format(v.name),
            escape(model[v].domain.str(value))
        )

    with open(file_name, 'w') as f:
        f.write(dot_printer.gen_dot(res_graph, [
            dot_printer.DataPrinter('___orig', print_orig)
        ] + [
            dot_printer.DataPrinter(v.name, print_result_builder(v))
            for v in var_set
        ]))

class AnalysisResults(object):
    """
    Contains the results of the abstract semantics analysis.
    """
    def __init__(self, cfg, semantics, trace_domain, vars_domain,
                 evaluator, orig_subp):
        self.cfg = cfg
        self.semantics = semantics
        self.trace_domain = trace_domain
        self.vars_domain = vars_domain
        self.evaluator = evaluator
        self.orig_subp = orig_subp

    def save_cfg_to_file(self, file_name):
        """
        Prints the control-flow graph as a DOT file to the given file name.
        """
        _save_cfg_to(file_name, self.cfg)

    def save_results_to_file(self, file_name):
        """
        Prints the resulting graph as a DOT file to the given file name.
        Displays the state of each variable at each program point.
        """
        _build_resulting_graph(
            file_name,
            self.cfg,
            self.semantics,
            self.trace_domain,
            self.evaluator.model
        )

    @staticmethod
    def _to_state(env):
        """
        Given an environment as a map from variable to value (domain element),
        generate a state vector as used by the expression evaluator. A state
        vector maps the INDEX of a variable to its value instead of the mapping
        the variable itself. For efficiency, we use a tuple instead of a dict
        to contain the mapping, although not all indices necessarily have an
        image. Those that don't have their component equal to None.

        :param dict[Variable, object] env: The environment, a map from
            variables to the domain element that they hold.

        :rtype: tuple[object]
        """

        # Because the set of indices may have holes, the maximal index is not
        # necessarily equal to the number of variables minus one. Therefore,
        # we first compute the appropriate length of the list that we will
        # use to store the state.

        max_index = (
            max(v.data.index for v in env.keys()) if len(env) > 0 else 0
        )
        state_vector = [None] * (max_index + 1)

        for var, value in env.iteritems():
            state_vector[var.data.index] = value

        return tuple(state_vector)

    def eval_at(self, node, expr):
        """
        Given a program point, evaluates for each program trace available at
        this program point the given expression.

        :param Digraph.Node node: The program point at which to evaluate the
            expression.

        :param irt.Expr expr: The expression to evaluate using the knowledge
            at that specific program point.

        :rtype: dict[frozenset[Digraph.Node], object]
        """
        return {
            trace: self.evaluator.eval(
                expr,
                self._to_state(env)
            )
            for trace, env in self.semantics[node].iteritems()
        }

_unit_domain = domains.Product()

def compute_semantics(prog, prog_model, merge_pred_builder, arg_values=None):
    evaluator = ExprEvaluator(prog_model)
    solver = ExprSolver(prog_model)

    # setup widening configuration
    visit_counter = KeyCounter()
    widening_delay = 5
    narrowing_delay = 3

    def do_widen(counter):
        # will widen when counter == widen_delay, then narrow. If it has not
        # converged after narrow_delay is reached, widening is triggered again
        # but without a follow-up narrowing.
        return (counter == widening_delay
                or counter >= narrowing_delay + widening_delay)

    cfg = prog.visit(CFGBuilder())
    roots = cfg.roots()
    non_roots = [n for n in cfg.nodes if n not in roots]

    # find the variables that appear in the program
    var_set = set(n for n in prog_model.keys() if isinstance(n, Variable))

    # build an index
    indexed_vars = {var.data.index: var for var in var_set}
    last_index = max(indexed_vars.keys()) if len(indexed_vars) > 0 else -1

    # define the variables domain
    vars_domain = domains.Product(*(
        prog_model[indexed_vars[i]].domain
        if i in indexed_vars else _unit_domain
        for i in range(last_index + 1)
    ))

    # define the trace domain
    trace_domain = _SimpleTraceLattice(cfg.nodes)

    # define the State domain that we track at each program point.
    lat = domains.Powerset(
        domains.Product(
            trace_domain,
            vars_domain
        ),
        merge_pred_builder.build(
            trace_domain,
            vars_domain
        ),
        None  # We don't need a top element here.
    )

    # the transfer function
    transfer_func = _VarTracker(var_set, vars_domain, evaluator, solver)

    def transfer(states, node, inputs):
        """
        Applies the transfer function to the given nodes using the abstract
        domain element that represents an over-approximation of the states at
        the program points that precede it.

        Returns a new element of the abstract domain, representing the state
        at this node after the application of the transfer function.

        :param dict[Digraph.Node, object] states: The state at each program
            point.
        :param Digraph.Node node: The node on which to apply the transfer
            function.
        :param object inputs: The element of the abstract domain that
            represents the state of all the predecessors combined.
        :rtype: object
        """

        transferred = (
            (
                trace,
                node.data.node.visit(transfer_func, values)
                if node.data.node is not None else values
            )
            for trace, values in inputs
        )

        output = lat.build([
            (
                trace_domain.join(trace, trace_domain.build([node])),
                values
            )
            for trace, values in transferred
            if not vars_domain.is_empty(values)
        ])

        if node.data.is_widening_point:
            if do_widen(visit_counter.get_incr(node)):
                output = lat.update(states[node], output, True)

        return output

    def iterate_once(states, ordering):
        """
        Perform one iteration over the system of data-flow equations associated
        with the given subset of nodes.

        This procedure updates "states" in-place.

        :param dict[Digraph.Node, object] states: The state at each program
            point.
        :param Digraph.HierarchicalOrdering ordering: The ordering, describing:
            - The order in which to apply the transfer functions to each node.
            - The subset of program points to consider in this iteration.
        """
        for elem, is_node in ordering:
            if is_node:
                states[elem] = transfer(states, elem, reduce(
                    lat.join,
                    (states[anc] for anc in cfg.ancestors(elem))
                ))
            else:
                fix(states, elem)

    def is_eq(last, current):
        """
        Given two dictionaries describing the state at each program point,
        returns True iff the state at each program point in `last` is equal
        to the state at *those* program points in `current`.

        :param dict[Digraph.Node, object] last: The last state (possibly
            containing less nodes than current).
        :param dict[Digraph.Node, object] current: The current state.
        :rtype: bool
        """
        for n, x in last.iteritems():
            if not lat.eq(x, current[n]):
                return False
        return True

    def sub_states(states, ordering):
        """
        Creates a copy of the given "states" dictionary which contains entries
        only for the nodes that are given by the first level of the ordering
        (i.e. from the current component).

        :type states: dict[Digraph.Node, object]
        :type ordering: Digraph.HierarchicalOrdering
        :rtype: dict[Digraph.Node, object]
        """
        return {
            elem: states[elem]
            for elem, is_node in ordering
            if is_node
        }

    def fix(states, ordering):
        """
        Solve the data-flow equations on the given subset of nodes of the CFG.
        Finds a fix-point by successive iterations.

        This procedure updates "states" in-place.

        :param dict[Digraph.Node, object] states: The state at each program
            point.
        :param Digraph.HierarchicalOrdering ordering: The ordering, describing:
            - The order in which to apply the transfer functions to each node.
            - The subset of program points for which to find a fix-point.
        """
        last = sub_states(states, ordering)
        iterate_once(states, ordering)

        # loop until the current state is equivalent to the last one.
        while not is_eq(last, states):
            last = sub_states(states, ordering)
            iterate_once(states, ordering)

    # initial state of the variables at the entry of the program
    init_vars = tuple(
        arg_values[indexed_vars[i]]
        if (i in indexed_vars and
            arg_values is not None and
            indexed_vars[i] in arg_values)
        else vars_domain.domains[i].top
        for i in range(last_index + 1)
    )

    # initial state to use at the entry of the program
    init_lat = lat.build([(trace_domain.bottom, init_vars)])

    # initial state at each program point the program
    states = concat_dicts(
        {n: transfer({}, n, init_lat) for n in roots},
        {n: lat.bottom for n in non_roots}
    )

    # Find a fix-point.
    fix(states, cfg.subgraph(non_roots).flat_topological_ordering())

    formatted_results = {
        node: {
            trace: {
                v: values[v.data.index] for v in var_set
            }
            for trace, values in state
        }
        for node, state in states.iteritems()
    }

    return AnalysisResults(
        cfg,
        formatted_results,
        trace_domain,
        vars_domain,
        evaluator,
        prog.data.fun_id
    )


import numpy as np

def compute_coefficients(type_exp):
    """
        This function computes the coefficients for the multicompartment model.
        Input: none
        Output: dict of coefficients. The main script reads the coeffs from it.
    """
    # init main dict
    main_dict = {}

    # Small script to compute the right coefficients.
    volume_brain_human = 1.0e6 #mm^3
    volume_brain_rat = 2311.
    mmHgToPa = 133.32

    # From El Bouri + Payne, we know
    #Kc = 4.28e-4 # in mm^3 s kg^-1
    Ka = 1.234
    Kv = 2.468

    Kc = 3.3e-3 # from smith et al (2014) "transmural variation and anisotropy..."

    # permeability and viscosity
    mu_blood = 2.67e-3 #Pa.s
    kc = Kc*1e-9*mu_blood
    ka = Ka*1e-9*mu_blood
    kv = Kv*1e-9*mu_blood

    # Resitance coeffs from Vinje (2020)
    mlTomm = 1000.0
    factor_resistance = mmHgToPa*60.0/mlTomm

    # Resistances in perivascular spaces
    R_pa = 1.14 * factor_resistance
    R_pv = 1.75e-3 * factor_resistance
    R_pc = 32.24 * factor_resistance

    # Resistance across membranes
    R_pa_e = 0.57 * factor_resistance
    R_pv_e = 0.64 * factor_resistance
    R_c_pc = 125.31 * factor_resistance

    ### Permeabilities
    # 1st step determine the ration
    # We know the permeability in the ECS and its resistance
    kappa_ECS = 2.0e-11 # mm^2
    R_ECS = 0.57 * factor_resistance
    mu_CSF = 7.e-4 # Pa.s
    ratio_L_A = kappa_ECS*R_ECS/mu_CSF
    #print('ratio L/A = ' +str(ratio_L_A))

    kappa_a = ka*1.e6
    kappa_v = kv*1.e6
    kappa_c = kc*1.e6

    kappa_pa = mu_CSF/R_pa*ratio_L_A
    kappa_pv = mu_CSF/R_pv*ratio_L_A
    kappa_pc = mu_CSF/R_pc*ratio_L_A

    #print("###  Permeabilities  ###\n")
    #print("kappa_a = " +str(kappa_a))
    #print("kappa_v = " +str(kappa_v))
    #print("kappa_c = " +str(kappa_c))

    #print("kappa_pa = " +str(kappa_pa))
    #print("kappa_pv = " +str(kappa_pv))
    #print("kappa_pc = " +str(kappa_pc))

    #print("kappa_ECS = " +str(kappa_ECS))

    ## Transfer coefficients
    # across BBB
    # Hydraulic conductivity
    L_a_e = 9.1e-10
    L_c_e = 1.0e-10
    L_v_e = 2.0e-11

    # surface to volume ration of vessels
    surface_v_to_volume = 3.  #mm^-1
    surface_a_to_volume = 3. #mm^-1
    surface_c_to_volume = 9.  #mm^-1

    # compute transfer from blood vessels to ECS
    gamma_a_e = L_a_e*surface_a_to_volume
    gamma_v_e = L_v_e*surface_v_to_volume
    gamma_c_e = L_c_e*surface_c_to_volume

    #print("\n###  Transfer  ###\n")
    #print("gamma_a_e = " +str(gamma_a_e))
    #print("gamma_v_e = " +str(gamma_v_e))
    #print("gamma_c_e = " +str(gamma_c_e))

    # PVS to tissue
    gamma_pa_e = 1.0/(R_pa_e*volume_brain_human)
    gamma_pv_e = 1.0/(R_pv_e*volume_brain_human)
    R_pc_e = (1/(gamma_c_e*volume_brain_human) - R_c_pc) # compute resistance
    gamma_pc_e = 1.0/(R_pc_e*volume_brain_human)

    #print("gamma_pa_e = " +str(gamma_pa_e))
    #print("gamma_pv_e = " +str(gamma_pv_e))
    #print("gamma_pc_e = " +str(gamma_pc_e))

    # From blood to PVS
    # Compute resistance from blood vessel to corresponding PVS
    R_a_pa = 1/(gamma_a_e*volume_brain_human)-R_pa_e
    R_v_pv = 1/(gamma_v_e*volume_brain_human)-R_pv_e

    gamma_a_pa = 1.0/(R_a_pa*volume_brain_human)
    gamma_v_pv = 1.0/(R_v_pv*volume_brain_human)
    gamma_c_pc = 1.0/(R_c_pc*volume_brain_human)

    #print("gamma_a_pa = " +str(gamma_a_pa))
    #print("gamma_v_pv = " +str(gamma_v_pv))
    #print("gamma_c_pc = " +str(gamma_c_pc))

    ## Transfer along vessels
    # From PVS arteries to PVS capillaries and PVS capillaries to PVS veins
    deltap_pa_pc = 1.0*mmHgToPa # difference of pressure between PVS arteries and PVS capillaries
    deltap_pc_pv = 0.25*mmHgToPa # difference of pressure between PVS arteries and PVS capillaries
    B_PVS = 3.38/60 # (mm^3/s) flow in PVS

    # transfer coefficients
    gamma_pa_pc = B_PVS/(deltap_pa_pc*volume_brain_rat)
    gamma_pc_pv = B_PVS/(deltap_pc_pv*volume_brain_rat)

    # From arteries to capillaries and capillaries to veins
    deltap_a_c = 40.0*mmHgToPa # difference of pressure between arteries and capillaries
    deltap_c_v = 13.0*mmHgToPa # difference of pressure between capillaries and veins
    B_blood = 116.*1.e3/100*2./60 # flow in blood network

    gamma_a_c = B_blood/(deltap_a_c*volume_brain_rat)
    gamma_c_v = B_blood/(deltap_c_v*volume_brain_rat)

    #print("\ngamma_a_c = " + str(gamma_a_c))
    #print("gamma_c_v = " +str(gamma_c_v))
    #print("gamma_pa_pc = " +str(gamma_pa_pc))
    #print("gamma_pc_pv = " +str(gamma_pc_pv))

    # Coefficient boundaries
    #print("\n### Hydraulic conductivities at pial surface")
    surface_brain = 1750*1.e2
    gamma_e_SAS = 1/(2*R_pa)

    #print("L_e_SAS = " + str(gamma_e_SAS/surface_brain))
    #print("L_pa_SAS = " + str(gamma_pa_e/surface_brain*volume_brain_human))
    L_e_SAS = gamma_e_SAS/surface_brain
    L_pa_SAS = gamma_pa_e/surface_brain*volume_brain_human

    ### Diffusive Permeabilities
    #print("\n### Computing diffusive permeability coefficients ###")
    #print("## For Inulin ##")
    #print("# Across the Astrocyte endfeet barrier# \n")
    # Computing diffusive permeabilities from Fu article  "A model for the blood–brain barrier permeability to water and small solutes"

    ## For Inulin
    # Computing effective diffusion in AEF pores at the capillary level
    Dv = 10000.0e-6 # diameter capillary
    D_Free_Inulin = 2.98e-4 # Free diffusion coefficient # mm²/s
    a_Inulin = 15.2e-7 # mm

    L_AEF = 1000.0e-6 # Width of AEF (assumed constant at the different levels)
    B_AEF = 25.0e-6 # mm can inccrease up to 1000e-6 (We assume very small because we are at the capillary level)
    beta = a_Inulin/B_AEF
    D_eff_AEF_Inulin = D_Free_Inulin*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))

    R_AEF_Inulin = L_AEF/(2*B_AEF*D_eff_AEF_Inulin)

    P_AEF_Inulin = 1/R_AEF_Inulin*1/(np.pi*Dv)
    lambda_pc_e_Inulin = P_AEF_Inulin*surface_c_to_volume

    #print("P_pc,e_Inulin = "  +str(P_AEF_Inulin))

    # Computing effective diffusion in AEF pores at the venous and arterial levels
    Dv = 50000.0e-6 # diameter arteriole or venule
    Dv = 38.e-3

    B_AEF = 250.0e-6 # mm can decrease to 2.5e-6 (We assume large because we are at the venous and arterial levels)
    beta = a_Inulin/B_AEF
    D_eff_AEF_Inulin = D_Free_Inulin*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))

    R_AEF_Inulin = L_AEF/(2*B_AEF*D_eff_AEF_Inulin)

    P_AEF_Inulin = 1/R_AEF_Inulin*1/(np.pi*Dv)
    lambda_pa_e_Inulin = P_AEF_Inulin*surface_a_to_volume
    lambda_pv_e_Inulin = P_AEF_Inulin*surface_v_to_volume
    #print("P_pa,e_Inulin = P_pv,e_Inulin ="  +str(P_AEF_Inulin))

    #print("\nlambda_pa,e_Inulin = "  +str(lambda_pa_e_Inulin))
    #print("lambda_pc,e_Inulin = "  +str(lambda_pc_e_Inulin))
    #print("lambda_pv,e_Inulin = "  +str(lambda_pv_e_Inulin))

    #print("\n## For Amyloid-beta ##")
    #print("# Across the BBB#") # We need to bring down all the different layers
    ## Capillary level
    # Astrocyte process
    L_AEF = 1000.0e-6 # Width of AEF (assumed constant at the different levels)
    B_AEF = 10.0e-6 # mm can inccrease up to 1000e-6 (We assume very small because we are at the capillary level)
    a_AB = 0.9e-6 # mm
    beta = a_AB/B_AEF
    D_Free_AB = 1.8e-6
    D_eff_AEF_AB = D_Free_AB*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))

    R_AEF_AB = L_AEF/(2*B_AEF*D_eff_AEF_AB)
    P_AEF_AB = 1/R_AEF_AB*1/(np.pi*Dv)

    # We already know
    lambda_pc_e_AB = P_AEF_AB*surface_c_to_volume

    # SGL layer
    B= 9.0e-6 # width of inter-endothelial cleft
    L_f = 250.0e-6 # Glycocalyx layer thickness
    r_f = 6e-6 # radius of fibers
    eps = 1-0.326 # volume fraction of void in glycocalyx layer
    D_f_AB =  D_Free_AB*(np.exp( -pow(1-eps,0.5)*(1+ a_AB/(r_f )) )) # unordered fibers see Michel 1999

    R_f = L_f/(2*B*D_f_AB)

    # inter-endothelial cleft
    L_1 = 350.0e-6 # Distance between the tight junction strand and the front of the inter-endothelial cleft

    beta = a_AB/B
    D_c_AB =  D_Free_AB*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))
    R_1_AB = L_1/(2*B*D_c_AB)

    # small silt in the junction strand
    L_jun = 11.0e-6 # Thickness of tight junction strand
    B_s = 2.5e-6  # Width of the small slit of the tight junction strand (Allt and Lawrenson, 1997; Cassella et al., 1997)
    beta = a_AB/B_s
    D_sl_AB =  D_Free_AB*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))
    R_2_AB = L_jun/(2*B_s*D_sl_AB)

    # small slit in the junction strand
    L = 700.0e-6 # Total length of the cleft region (Schulze and Firth, 1992)
    R_3_AB = (L-L_1-L_f)/(2*B*D_c_AB)

    # Basement membrane (fibrous membrane)
    W_a = 2500.0e-6 # Length of astrocyte foot processes (Farkas and Luiten, 2001)
    r_b = 2.7e-6 # radius of BM fibers Michel (1999)
    eps_b = 1-0.5 # void volume fraction in BM Michel (1999)
    L_b = 30.0e-6 # Thickness of the basement membrane (Paulson and Newman, 1987; Farkas and Luiten, 2001)
    D_b_AB =  D_Free_AB*(np.exp( -pow(1-eps_b,0.5)*(1+ a_AB/(r_b )) )) # unordered fibers see Michel 1999
    R_4_AB = W_a/(4*L_b*D_b_AB)

    # Permeability across total BBB at the capillary level
    Dv = 10.0e-3
    P_BBB_AB = 1/(R_AEF_AB + R_f  + R_1_AB + R_2_AB + R_3_AB + R_4_AB) * 1/(np.pi * Dv)

    lambda_c_e_AB = P_BBB_AB * surface_c_to_volume
    P_BBB_del_AEF_AB  = 1/(R_f  + R_1_AB + R_2_AB + R_3_AB + R_4_AB) * 1/(np.pi * Dv)
    lambda_c_pc_AB = P_BBB_del_AEF_AB * surface_c_to_volume

    ### Now the arterial space, we assume that the junctions are tight because the BBB is known to be less permeable at this level
    # Astrocyte process
    L_AEF = 1000.0e-6 # Width of AEF (assumed constant at the different levels)
    B_AEF = 500.0e-6 # mm can inccrease up to 1000e-6 (We assume large because we are at the artery level)
    beta = a_AB/B_AEF
    D_eff_AEF_AB = D_Free_AB*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))

    R_AEF_AB = L_AEF/(2*B_AEF*D_eff_AEF_AB)
    P_AEF_AB = 1/R_AEF_AB*1/(np.pi*Dv)

    # We already know
    lambda_pa_e_AB = P_AEF_AB*surface_a_to_volume

    # SGL layer
    B= 9.0e-6 # width of inter-endothelial cleft
    L_f = 400.0e-6 # Glycocalyx layer thickness (assumed to be large at this level)
    r_f = 6e-6 # radius of fibers
    eps = 1-0.326 # volume fraction of void in glycocalyx layer
    D_f_AB =  D_Free_AB*(np.exp( -pow(1-eps,0.5)*(1+ a_AB/(r_f )) )) # unordered fibers see Michel 1999

    R_f = L_f/(2*B*D_f_AB)

    # inter-endothelial cleft
    L_1 = 350.0e-6 # Distance between the tight junction strand and the front of the inter-endothelial cleft

    beta = a_AB/B
    D_c_AB =  D_Free_AB*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))
    R_1_AB = L_1/(2*B*D_c_AB)

    # small silt in the junction strand
    L_jun = 11.0e-6 # Thickness of tight junction strand
    B_s = 0.5e-6  # Width of the small slit of the tight junction strand (Allt and Lawrenson, 1997; Cassella et al., 1997)
    beta = a_AB/B_s
    D_sl_AB =  D_Free_AB*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))
    R_2_AB = L_jun/(2*B_s*D_sl_AB)

    # small slit in the junction strand
    L = 700.0e-6 # Total length of the cleft region (Schulze and Firth, 1992)
    R_3_AB = (L-L_1-L_f)/(2*B*D_c_AB)

    # Basement membrane (fibrous membrane)
    W_a = 2500.0e-6 # Length of astrocyte foot processes (Farkas and Luiten, 2001)
    r_b = 2.7e-6 # radius of BM fibers Michel (1999)
    eps_b = 1-0.5 # void volume fraction in BM Michel (1999)
    L_b = 80.0e-6 # Thickness of the basement membrane (Paulson and Newman, 1987; Farkas and Luiten, 2001)
    D_b_AB =  D_Free_AB*(np.exp( -pow(1-eps_b,0.5)*(1+ a_AB/(r_b )) )) # unordered fibers see Michel 1999
    R_4_AB = W_a/(4*L_b*D_b_AB)

    # Permeability across total BBB at the capillary level
    Dv = 50.0e-3
    P_BBB_AB = 1/(R_AEF_AB + R_f  + R_1_AB + R_2_AB + R_3_AB + R_4_AB) * 1/(np.pi * Dv)

    lambda_a_e_AB = P_BBB_AB * surface_a_to_volume

    P_BBB_del_AEF_AB = 1/( R_f  + R_1_AB + R_2_AB + R_3_AB + R_4_AB) * 1/(np.pi * Dv)

    lambda_a_pa_AB = P_BBB_del_AEF_AB * surface_a_to_volume

    ### Now the VENOUS space, we assume that the junctions are loose because the BBB is known to be more permeable at this level (less claudin expression)
    # Astrocyte process
    L_AEF = 1000.0e-6 # Width of AEF (assumed constant at the different levels)
    B_AEF = 500.0e-6 # mm can inccrease up to 1000e-6 (We assume large because we are at the artery level)
    beta = a_AB/B_AEF
    D_eff_AEF_AB = D_Free_AB*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))

    R_AEF_AB = L_AEF/(2*B_AEF*D_eff_AEF_AB)
    P_AEF_AB = 1/R_AEF_AB*1/(np.pi*Dv)

    # We already know
    lambda_pv_e_AB = P_AEF_AB*surface_v_to_volume

    # SGL layer
    B= 9.0e-6 # width of inter-endothelial cleft
    L_f = 100.0e-6 # Glycocalyx layer thickness (assumed to be large at this level)
    r_f = 6e-6 # radius of fibers
    eps = 1-0.326 # volume fraction of void in glycocalyx layer
    D_f_AB =  D_Free_AB*(np.exp( -pow(1-eps,0.5)*(1+ a_AB/(r_f )) )) # unordered fibers see Michel 1999

    R_f = L_f/(2*B*D_f_AB)

    # inter-endothelial cleft
    L_1 = 350.0e-6 # Distance between the tight junction strand and the front of the inter-endothelial cleft

    beta = a_AB/B
    D_c_AB =  D_Free_AB*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))
    R_1_AB = L_1/(2*B*D_c_AB)

    # small silt in the junction strand
    L_jun = 11.0e-6 # Thickness of tight junction strand
    B_s = 10.0e-6  # Width of the small slit of the tight junction strand (Allt and Lawrenson, 1997; Cassella et al., 1997)
    beta = a_AB/B_s
    D_sl_AB =  D_Free_AB*(1-2.10444*pow(beta,6) +2.08877*pow(beta,3) - 0.094813*pow(beta,5) - 1.372*pow(beta,6))
    R_2_AB = L_jun/(2*B_s*D_sl_AB)

    # small slit in the junction strand
    L = 700.0e-6 # Total length of the cleft region (Schulze and Firth, 1992)
    R_3_AB = (L-L_1-L_f)/(2*B*D_c_AB)

    # Basement membrane (fibrous membrane)
    W_a = 2500.0e-6 # Length of astrocyte foot processes (Farkas and Luiten, 2001)
    r_b = 2.7e-6 # radius of BM fibers Michel (1999)
    eps_b = 1-0.5 # void volume fraction in BM Michel (1999)
    L_b = 20.0e-6 # Thickness of the basement membrane (Paulson and Newman, 1987; Farkas and Luiten, 2001)
    D_b_AB =  D_Free_AB*(np.exp( -pow(1-eps_b,0.5)*(1+ a_AB/(r_b )) )) # unordered fibers see Michel 1999
    R_4_AB = W_a/(4*L_b*D_b_AB)

    # Permeability across total BBB at the capillary level
    Dv = 50.0e-3
    P_BBB_AB = 1/(R_AEF_AB + R_f  + R_1_AB + R_2_AB + R_3_AB + R_4_AB) * 1/(np.pi * Dv)

    lambda_v_e_AB = P_BBB_AB * surface_v_to_volume
    P_BBB_del_AEF_AB = 1/(R_f  + R_1_AB + R_2_AB + R_3_AB + R_4_AB) * 1/(np.pi * Dv)
    lambda_v_pv_AB = P_BBB_del_AEF_AB * surface_v_to_volume

    #print("\nlambda_a_e_AB = "+str(lambda_a_e_AB))
    #print("lambda_v_e_AB = "+str(lambda_v_e_AB))
    #print("lambda_c_e_AB = "+str(lambda_c_e_AB))

    #print("\n## Across the AEF")

    #print("\nlambda_pa_e_AB = "+str(lambda_pa_e_AB))
    #print("lambda_pv_e_AB = "+str(lambda_pv_e_AB))
    #print("lambda_pc_e_AB = "+str(lambda_pc_e_AB))

    #print("\n## Across the BBB without the AEF")
    #print("\nlambda_a_pa_AB = "+str(lambda_a_pa_AB))
    #print("lambda_v_pv_AB = "+str(lambda_v_pv_AB))
    #print("lambda_c_pc_AB = "+str(lambda_c_pc_AB))

    ### Convective mass exchange
    ## Note The reflection coefficient is often found from microvessels but it depends on the membrane itself.
    ## Whether this value changes from the AEF membrane only is unclear. Might need further investigation.
    ## Inulin
    sigma_reflect = 0.2

    tilde_gamma_pa_e_Inulin = gamma_pa_e*(1-sigma_reflect)
    tilde_gamma_pv_e_Inulin = gamma_pv_e*(1-sigma_reflect)
    tilde_gamma_pc_e_Inulin = gamma_pc_e*(1-sigma_reflect)

    # For connected channels, we assume that there is not a membrane (so no diffusive transfer)
    sigma_reflect = 0.0

    tilde_gamma_pa_pc_Inulin = gamma_pa_pc
    tilde_gamma_pc_pv_Inulin = gamma_pc_pv

    ## Amyloid beta
    sigma_reflect = 0.1

    tilde_gamma_a_e_AB = gamma_a_e*(1-sigma_reflect)
    tilde_gamma_v_e_AB = gamma_v_e*(1-sigma_reflect)
    tilde_gamma_c_e_AB = gamma_c_e*(1-sigma_reflect)

    tilde_gamma_pa_e_AB = gamma_pa_e*(1-sigma_reflect)
    tilde_gamma_pv_e_AB = gamma_pv_e*(1-sigma_reflect)
    tilde_gamma_pc_e_AB = gamma_pc_e*(1-sigma_reflect)

    tilde_gamma_a_pa_AB = gamma_a_pa*(1-sigma_reflect)
    tilde_gamma_v_pv_AB = gamma_v_pv*(1-sigma_reflect)
    tilde_gamma_c_pc_AB = gamma_c_pc*(1-sigma_reflect)

    # For connected channels, we assume that there is not a membrane (so no diffusive transfer)
    sigma_reflect = 0.0

    tilde_gamma_a_c_AB = gamma_a_c
    tilde_gamma_c_v_AB = gamma_c_v
    tilde_gamma_pa_pc_AB = gamma_pa_pc
    tilde_gamma_pc_pv_AB = gamma_pc_pv

    #print("\n### Convective mass transfer ###\n")
    #print("## For Inulin\n")
    #print("tilde_gamma_pa_e_Inulin = "+str(tilde_gamma_pa_e_Inulin))
    #print("tilde_gamma_pv_e_Inulin = "+str(tilde_gamma_pv_e_Inulin))
    #print("tilde_gamma_pc_e_Inulin = "+str(tilde_gamma_pc_e_Inulin))

    #print("\ntilde_gamma_pa_pc_Inulin = "+str(tilde_gamma_pa_pc_Inulin))
    #print("tilde_gamma_pc_pv_Inulin = "+str(tilde_gamma_pc_pv_Inulin))

    #print("\n## For Amyloid-Beta\n")

    #print("tilde_gamma_a_e_AB = " + str(tilde_gamma_a_e_AB))
    #print("tilde_gamma_v_e_AB = " + str(tilde_gamma_v_e_AB))
    #print("tilde_gamma_c_e_AB  = "+ str(tilde_gamma_c_e_AB))

    #print("tilde_gamma_pa_e_AB = " + str(tilde_gamma_pa_e_AB))
    #print("tilde_gamma_pv_e_AB = " + str(tilde_gamma_pv_e_AB))
    #print("tilde_gamma_pc_e_AB = " + str(tilde_gamma_pc_e_AB))

    #print("tilde_gamma_a_pa_AB = " + str(tilde_gamma_a_pa_AB))
    #print("tilde_gamma_v_pv_AB = " + str(tilde_gamma_v_pv_AB))
    #print("tilde_gamma_c_pc_AB = " + str(tilde_gamma_c_pc_AB))

    # Save to dictionary

    volume_fraction_of_blood = 3.29/100.
    arteries_frac = 0.21
    veins_frac = 0.46
    capillaries_frac = 0.33
    PVS_fraction = 1./100.

    main_dict["phi_e"] = 0.14
    main_dict["phi_a"] = volume_fraction_of_blood*arteries_frac
    main_dict["phi_v"] = volume_fraction_of_blood*veins_frac
    main_dict["phi_c"] = volume_fraction_of_blood*capillaries_frac
    main_dict["phi_pa"] = PVS_fraction*arteries_frac
    main_dict["phi_pv"] = PVS_fraction*veins_frac
    main_dict["phi_pc"] = PVS_fraction*capillaries_frac

    #print("\n### Baseline porosities ###\n")
    #print("phi_e = "+str(main_dict["phi_e"]))
    #print("phi_a = "+str(main_dict["phi_a"]))
    #print("phi_v = "+str(main_dict["phi_v"]))
    #print("phi_c = "+str(main_dict["phi_c"]))
    #print("phi_pa = "+str(main_dict["phi_pa"]))
    #print("phi_pv = "+str(main_dict["phi_pv"]))
    #print("phi_pc = "+str(main_dict["phi_pc"]))

    main_dict["mu_blood"] = mu_blood
    main_dict["mu_csf"] = mu_CSF

    main_dict["gamma_a_e"] = gamma_a_e
    main_dict["gamma_v_e"] = gamma_v_e
    main_dict["gamma_c_e"] = gamma_c_e
    main_dict["gamma_pa_e"] = gamma_pa_e
    main_dict["gamma_pv_e"] = gamma_pv_e
    main_dict["gamma_pc_e"] = gamma_pc_e
    main_dict["gamma_a_pa"] = gamma_a_pa
    main_dict["gamma_v_pv"] = gamma_v_pv
    main_dict["gamma_c_pc"] = gamma_c_pc
    main_dict["gamma_a_c"] = gamma_a_c
    main_dict["gamma_c_v"] = gamma_c_v
    main_dict["gamma_pa_pc"] = gamma_pa_pc
    main_dict["gamma_pc_pv"] = gamma_pc_pv
    main_dict["kappa_a"] = kappa_a
    main_dict["kappa_v"] = kappa_v
    main_dict["kappa_c"] = kappa_c
    main_dict["kappa_pa"] = kappa_pa
    main_dict["kappa_pv"] = kappa_pv
    main_dict["kappa_pc"] = kappa_pc
    main_dict["kappa_e"] = kappa_ECS
    main_dict["lambda_pa_e_Inulin"] = lambda_pa_e_Inulin
    main_dict["lambda_pv_e_Inulin"] = lambda_pv_e_Inulin
    main_dict["lambda_pc_e_Inulin"] = lambda_pc_e_Inulin
    main_dict["tilde_gamma_pa_e_Inulin"] = tilde_gamma_pa_e_Inulin
    main_dict["tilde_gamma_pc_e_Inulin"] = tilde_gamma_pc_e_Inulin
    main_dict["tilde_gamma_pv_e_Inulin"] = tilde_gamma_pv_e_Inulin
    main_dict["tilde_gamma_pa_pc_Inulin"] = tilde_gamma_pa_pc_Inulin
    main_dict["tilde_gamma_pc_pv_Inulin"] = tilde_gamma_pc_pv_Inulin

    main_dict["L_e_SAS"] = L_e_SAS
    main_dict["L_pa_SAS"] = L_pa_SAS

    main_dict["blood_flow"] = B_blood

    if type_exp == "ECS enhanced":
        main_dict["phi_e"] = 0.23
        main_dict["kappa_e"] = kappa_ECS*5.5
    elif type_exp == "ECS+PVS enhanced":
        main_dict["phi_e"] = 0.23
        main_dict["kappa_e"] = kappa_ECS*5.5
        main_dict["phi_pa"] = main_dict["phi_pa"]*4.0
        main_dict["phi_pv"] = main_dict["phi_pv"]*4.0
        main_dict["phi_pc"] = main_dict["phi_pc"]*4.0
        main_dict["kappa_pa"] = kappa_pa*16.
        main_dict["kappa_pv"] = kappa_pv*16.
        main_dict["kappa_pc"] = kappa_pc *16.
    elif type_exp == "PVS only enhanced":
        main_dict["phi_pa"] = main_dict["phi_pa"]*4.0
        main_dict["phi_pv"] = main_dict["phi_pv"]*4.0
        main_dict["phi_pc"] = main_dict["phi_pc"]*4.0
        main_dict["kappa_pa"] = kappa_pa*16.
        main_dict["kappa_pv"] = kappa_pv*16.
        main_dict["kappa_pc"] = kappa_pc *16.

    return main_dict


# coding=utf-8
# Author:jonathon woo
# email:live.wujianxuan@gmail.com
# version:1.0.0
import json
import os
import hou

def loadCurve(*args, **kwargs):
    node = hou.selectedNodes()[0]
    node_name = node.name()
    test_path = r"D:\%s" % node_name
    if not os.path.exists(test_path):
        return
    # delete all child node
    for node in node.children():
        node.destroy()

    # read data
    for curve_data in os.listdir(test_path):
        full_data_path = os.path.join(test_path, curve_data)
        curve_name = os.path.splitext(curve_data)[0]

        # create curve node
        target_node = node.createNode("curve", curve_name)

        parm_dic = {}
        with open(full_data_path, "r") as outfile:
            all_info = outfile.read()
            parm_dic = json.loads(all_info)

        # add width parm to interface
        current_tmpparm = target_node.parmTemplateGroup()
        current_tmpparm.append(hou.FloatParmTemplate("width", "Width", 1))
        target_node.setParmTemplateGroup(current_tmpparm)
        target_node.setParms(parm_dic)
    print("Loaded!")


import logging
import cleaner
import renderer
import querier
import scraper
import cgi
from django.utils import simplejson
from google.appengine.ext import webapp
from google.appengine.ext.webapp.util import run_wsgi_app
from google.appengine.ext import db

"""Load custom Django template filters"""
webapp.template.register_template_library('djangofilters')

class Restaurant(db.Model):
    name = db.StringProperty()
    rating = db.FloatProperty()
    date = db.DateTimeProperty(auto_now_add=True)
    yelp_link = db.StringProperty()
    yelp_id = db.StringProperty()
    ot_link = db.StringProperty()
    ot_genre = db.StringProperty()
    ot_neighborhood = db.StringProperty()
    ot_mf = db.StringProperty()
    ot_sun = db.StringProperty()
    genre = db.StringProperty()
    neighborhood = db.StringProperty()
    address = db.StringProperty()
    lat = db.FloatProperty()
    lng = db.FloatProperty()
    g_lat = db.FloatProperty()
    g_lng = db.FloatProperty()

class MainPage(webapp.RequestHandler):
    def get(self):
        restaurants_query = Restaurant.all().order('-date')
        template_values = {
            'restaurants': restaurants_query
        }
        renderer.render(self, 'index.html', template_values)

class Adder(webapp.RequestHandler):
    def get(self):
        self.addAll()

    def addAll(self):
        ct = 0
        cap = 5
        for rest in scraper.get_all():
           ct += RestaurantMaker().make(rest)
           if ct > cap:
               self.redirect('/')
               return

class RestaurantSearchHandler(webapp.RequestHandler):
    def post(self):
        name = self.request.get('name')
        q = querier.Querier()
        yelp_result = q.search_yelp(name)
        self.response.out.write(simplejson.dumps( {
                'opentable':q.get_opentable_link(name),
                'yelp':yelp_result
                }))

class RestaurantPostHandler(webapp.RequestHandler):
    def post(self):
        RestaurantMaker().make(self.request.get('name'))

class RestaurantMaker:

    def getMeals(self, day, dict):
        if day in dict:
            if 'lunch' in dict[day] and 'dinner' in dict[day]:
                return 'Lunch and Dinner'
            elif 'lunch' in dict[day]:
                return 'Lunch'
            if 'dinner' in dict[day]:
                return 'Dinner'
        return ''

    def make(self, rest_data):
        r_name = rest_data['name']
        q = querier.Querier()
        data = {}
        data['ot_sun'] = self.getMeals('sunday', rest_data)
        data['ot_mf'] = self.getMeals('monday', rest_data)
        r_name_utf = r_name.decode('utf-8')
        rest = db.GqlQuery("SELECT * FROM Restaurant WHERE name = :name", name = r_name_utf).get()
        #don't store duplicates
        if rest is not None:
            return 0
        yelp_r = q.search_yelp(r_name)
        if yelp_r is None:
            return 0
        logging.info(r_name)
        data['name'] = r_name
        data['ot_genre'] = rest_data['cuisine']
        data['ot_neighborhood'] = rest_data['neighborhood']
        data['ot_link'] = q.get_opentable_link(r_name)
        data['yelp_link'] = yelp_r[0]['url']
        data['yelp_id'] = yelp_r[0]['id']
        data['rating'] = yelp_r[0]['avg_rating']
        data['lat'] = yelp_r[0]['latitude']
        data['lng'] = yelp_r[0]['longitude']
        data['phone'] = yelp_r[0]['phone']
        data['address'] = self.stringify(yelp_r[0], ['address1', 'address2', 'address3', 'zip'])
        data['genre'] = self.stringify(yelp_r[0]['categories'], ['name'])
        data['neighborhood'] = self.stringify(yelp_r[0]['neighborhoods'], ['name'])
        data['address'] = yelp_r[0]['address1']
        self.add(data)
        return 1

    def stringify(self, hash, keys):
        s = []
        if len(keys) == 1:
            for sub in hash:
                if sub[keys[0]] and sub[keys[0]] != '':
                    s.append(sub[keys[0]])
        else:
            for key in keys:
                if hash[key] and hash[key] != '':
                    s.append(hash[key])
        s.sort()
        return ', '.join(s)

    def add(self, data):
        utf_name = data['name'].decode('utf-8')
        rest = db.GqlQuery("SELECT * FROM Restaurant WHERE name = :name", name = utf_name).get()
        if rest is None:
            rest = Restaurant()
        #rest.name = data['name']
        rest.name = utf_name
        rest.ot_link = data['ot_link']
        rest.yelp_link = data['yelp_link']
        rest.rating = data['rating']
        rest.neighborhood = data['neighborhood']
        rest.genre = data['genre']
        rest.lat = data['lat']
        rest.lng = data['lng']
        rest.address = data['address']
        rest.ot_sun = data['ot_sun']
        rest.ot_mf = data['ot_mf']
        rest.yelp_id = data['yelp_id']
        rest.ot_genre = data['ot_genre']
        rest.ot_neighborhood = data['ot_neighborhood']
        rest.put()

application = webapp.WSGIApplication(
                                     [('/', MainPage),
                                      #('/findrestaurant', RestaurantSearchHandler),
                                      #('/addrestaurant', RestaurantPostHandler),
                                      ('/cleanall', cleaner.DataCleaner),
                                      ('/addall', Adder)],
                                     debug=True)
def main():
    run_wsgi_app(application)

if __name__ == "__main__":
    main()


debug = False
def readFile(filename):
    with open(filename, 'r') as myfile:
        # First line is the smax
        text = myfile.readline()[:-1]#
        i=1
        for line in myfile:
            # Read the line in a list
            line = line[:-1]
            if debug:
                print('-'*20)
                print(line)
            [B, N] = [int(e) for e in line.split(' ')]
            # Give the line without \n
            line = myfile.readline()[:-1]
            M = [int(e) for e in line.split(' ')]
            if debug:
                print(M)
            # solve the problem
            res = solve(B, N, M)
            print("Case #{}: {}".format(i, res))
            i += 1
            if debug:
                if i==10:
                    return 0
    return 0

def lcm(x, y):
    """This function takes two
        integers and returns the L.C.M."""
    # choose the greater number
    if( x > y):
        greater = x
    else:
        greater = y

    while(True):
        if((greater % x == 0) and (greater % y == 0)):
            lcm = greater
            break
        greater += 1
    return lcm

def lcm_l(L):
    res = L[0]
    for e in L[1:]:
        res = lcm(res,e)
    return res

def solve(B, N, M):
    global debug
    # reduce the problem
    c = lcm_l(M)
    d = sum([c//e for e in M])

    if debug:
        print("ppcm", c, "|||| clients par ppcm", d)

    if N > d:
        i = N - (N % d)
        if i == N:
            mx =  max(M)
            mxi = 0
            for j,e in enumerate(M):
                if e <= mx:
                    mx = e
                    mxi = j
            return mxi+1

    else:
        i = 0
    i+=1

    if debug:
        print("start :", i)

    ttl = []
    for e in range(B):
        ttl.append(0)
    res = 1
    clients = [0]*len(M)
    while i < N+1:
        free = False
        # take care of consumer i
        for j, t in enumerate(ttl):
        # free barber
            if t==0:
                free = True
                clients[j] = i
                ttl[j] += M[j]
                res = j + 1
                i += 1
                if debug:
                    print(clients)
                break
        if not free:
            for j in range(len(ttl)):
                ttl[j] -= 1
    return res

if __name__ == "__main__":
    filename = "test.txt"
    filename = "B-small-attempt6.in"
    readFile(filename)


#!/usr/bin/env python3

"""Rod cutting problem, with memoization."""

from typing import Sequence

# Prices for a rod of size i (for rod sizes [0,10]).
#.        0. 1  2. 3. 4. 5.  6.  7.  8.  9,  10
PRICES = [0, 1, 5, 8, 9, 10, 17, 17, 20, 24, 30]

# Cache results from find_best_cut
FIND_BEST_CUT_MEMO = {}

# Returns optimal cut strategy, in lengths, for a rod of input size k
def find_best_cut(k : int) -> Sequence[int]:
    if k == 0:
        return []

    if k in FIND_BEST_CUT_MEMO:
        return FIND_BEST_CUT_MEMO[k]

    best_value = -1
    best_cut = []

    # Start with full length log
    if k < len(PRICES):
        best_value = PRICES[k]
        best_cut = []

    for first_cut in range(1, min(len(PRICES), k)):
        remaining_len = k - first_cut
        smaller_cut = find_best_cut(remaining_len)
        curr_value = PRICES[first_cut] + _cut_value(remaining_len, smaller_cut)

        if curr_value > best_value:
            best_value = curr_value
            best_cut = [ first_cut ]
            best_cut.extend(smaller_cut)

    FIND_BEST_CUT_MEMO[k] = best_cut
    return best_cut

def _cut_value(k : int, cut_lengths : Sequence[int]) -> int:
    """Return the value of a cut rod of input size "k", with the given cut sequence."""
    result = 0
    total_len = 0
    for cut in cut_lengths:
        result += PRICES[cut]
        total_len += cut

    remainder = k - total_len
    if remainder > 0:
        result += PRICES[remainder]

    assert remainder >= 0, f'Expected length: {k}, computed {total_len}'
    return result

def print_cut_rod(k):
    cut_lengths = find_best_cut(k)

    print(f'Rod length: {k}, value: ' + str(_cut_value(k, cut_lengths)))
    print('Cuts lengths: [' + ', '.join([str(c) for c in cut_lengths]) + ']')

    line = ''
    total_len = 0
    for cut in cut_lengths:
        total_len += cut
        line += '. ' * (cut - 1)
        line += '.'
        line += '|'
    line += '. ' * (k - total_len) # remainder

    print(line)
    print()

print_cut_rod(1)
print_cut_rod(2)
print_cut_rod(7)
print_cut_rod(10)
print_cut_rod(11)
print_cut_rod(25)


import numpy as np
import sys

UNIT_SCALE=1e-03 ##Unit scales = milli
TIME_UNIT=UNIT_SCALE
PARAM_SCALING=1/UNIT_SCALE #Constants for Time, Conductances, Capacitances are up scaled for computation, Potentials remain same

#lists for [a,d] :amplitudes, decay times pairs and, at the last, rho: common rise times
AMPA_d_dlt=np.array([0.3351*1e-03, 1.651*1e-03])                *PARAM_SCALING  #Time terms are upscaled by milli
AMPA_d_rho= 0.3274*1e-03                                        *PARAM_SCALING
AMPA_s_dlt=np.array([0.4*1e-03, 4.899*1e-03, 43.1*1e-03])       *PARAM_SCALING
AMPA_s_rho= 0.5548*1e-03                                        *PARAM_SCALING
NMDA_dlt=np.array([13.52*1e-03, 121.9*1e-03])                   *PARAM_SCALING
NMDA_rho=0.8647*1e-03                                           *PARAM_SCALING
GABA_dlt=np.array([13.52*1e-03, 121.9*1e-03])                   *PARAM_SCALING
GABA_rho=0.8647*1e-03                                           *PARAM_SCALING

AMPA_d_amp=np.array([3.724*1e-09, 0.3033*1e-09])                *PARAM_SCALING*1  #Amp & conductances upscaled by milli
AMPA_s_amp=np.array([0.2487*1e-09,0.2799*1e-09, 0.1268*1e-09])  *PARAM_SCALING*1
NMDA_amp=np.array([17*1e-09, 2.645*1e-09])                      *PARAM_SCALING*1
GABA_amp=np.array([0.1*1e-09, 0.2*1e-09])                       *PARAM_SCALING*1

membrane_capacitance=3.22*1e-12                 *PARAM_SCALING #C_m
leak_conductance=1.06*1e-09                     *PARAM_SCALING #G_m
GABA_conductance=0.438*1e-09                    *PARAM_SCALING #G_GABAR (tonic)

leak_reversal_potential=-79.9*1e-03             #*PARAM_SCALING #E_m
GABA_reversal_potential=-79.1*1e-03             #*PARAM_SCALING #E_GABAR (tonic)
AMPA_reversal_potential=0                       #*PARAM_SCALING #E_AMPAR
NMDA_reversal_potential=0                       #*PARAM_SCALING #E_NMDAR
reset_membrane_potential=-63*1e-03              #*PARAM_SCALING #V_r
activation_threshold=-40*1e-03                  #*PARAM_SCALING #V_t

refractory_interval=2*1e-03                     *PARAM_SCALING #tau

r_AMPA_d=0.1249
r_AMPA_s=0.2792
r_NMDA  =0.0322
r_GABA =0.02

class Spiking_cells:
    def __init__(self,num_dend=1):
        #each refers to Synaptic weghts of each of the components
        self.p_AMPA_d_U=[]
        self.p_AMPA_d_R=[]
        self.p_AMPA_s_U=[]
        self.p_AMPA_s_R=[]
        self.p_NMDA_U  =[]
        self.p_NMDA_R  =[]

        self.p_GABA_U  =[]
        self.p_GABA_R  =[]
        for dend in range(num_dend):
            self.p_AMPA_d_U.append(r_AMPA_d)
            self.p_AMPA_d_R.append(1)
            self.p_AMPA_s_U.append(r_AMPA_s)
            self.p_AMPA_s_R.append(1)
            self.p_NMDA_U  .append(r_NMDA)
            self.p_NMDA_R  .append(1)

            self.p_GABA_U  .append(r_GABA) #GABA added
            self.p_GABA_R  .append(1)

        self.num_dend=num_dend
        self.membrane_potential= reset_membrane_potential
        self.MEM_POT_NOSTP= reset_membrane_potential

    def p_AMPA_d(self, ind_dend): return self.p_AMPA_d_U[ind_dend]*self.p_AMPA_d_R[ind_dend]
    def p_AMPA_s(self, ind_dend): return self.p_AMPA_s_U[ind_dend]*self.p_AMPA_s_R[ind_dend]
    def p_NMDA(self,   ind_dend): return self.p_NMDA_U[ind_dend]*self.p_NMDA_R[ind_dend]
    def p_GABA(self,   ind_dend): return self.p_GABA_U[ind_dend]*self.p_GABA_R[ind_dend] #GABA added

    def PLASTICITY(self, ind_dend, SPIKE):
        p_AMPA_d_U=self.p_AMPA_d_U[ind_dend]
        p_AMPA_s_U=self.p_AMPA_s_U[ind_dend]
        p_NMDA_U  =self.p_NMDA_U  [ind_dend]
        p_AMPA_d_R=self.p_AMPA_d_R[ind_dend]
        p_AMPA_s_R=self.p_AMPA_s_R[ind_dend]
        p_NMDA_R  =self.p_NMDA_R  [ind_dend]

        p_GABA_U  =self.p_GABA_U  [ind_dend] #GABA added
        p_GABA_R  =self.p_GABA_R  [ind_dend]
        if SPIKE:
            self.p_AMPA_d_U[ind_dend]+=r_AMPA_d*(1-p_AMPA_d_U)
            self.p_AMPA_s_U[ind_dend]+=r_AMPA_s*(1-p_AMPA_s_U)
            self.p_NMDA_U  [ind_dend]+=r_NMDA  *(1-p_NMDA_U  )
            self.p_AMPA_d_R[ind_dend]=p_AMPA_d_R*(1-p_AMPA_d_U)
            self.p_AMPA_s_R[ind_dend]=p_AMPA_s_R*(1-p_AMPA_s_U)
            self.p_NMDA_R  [ind_dend]=p_NMDA_R  *(1-p_NMDA_U  )
            self.p_GABA_U  [ind_dend]+=r_GABA   *(1-p_GABA_U  )  #GABA added
            self.p_GABA_R  [ind_dend]= p_GABA_R  *(1-p_GABA_U  )
        else:
            self.p_AMPA_d_U[ind_dend]-=(p_AMPA_d_U-r_AMPA_d)/6.394 #a Potentiation Recovery time constant
            self.p_AMPA_s_U[ind_dend]-=(p_AMPA_s_U-r_AMPA_s)/6.394 # For NMDA, But since no info. for AMPA so.
            self.p_NMDA_U  [ind_dend]-=(p_NMDA_U  -r_NMDA  )/6.394 # All for same
            self.p_AMPA_d_R[ind_dend]-=(p_AMPA_d_R-1)/131
            self.p_AMPA_s_R[ind_dend]-=(p_AMPA_s_R-1)/14.85
            self.p_NMDA_R  [ind_dend]-=(p_NMDA_R  -1)/236.1
            self.p_GABA_U  [ind_dend]-=(p_GABA_U-r_GABA   )/6.394 #GABA added
            self.p_GABA_R  [ind_dend]-=(p_GABA_R -1)/100

    def AMPA_conductance_evaluation(self, t, index_dend):
        d_wave=np.zeros(len(t))

        for ind, _ in enumerate(AMPA_d_amp):
            #d_wave+=AMPA_d_amp[ind]*(np.exp(-t/AMPA_d_rho)-np.exp(-t/AMPA_d_dlt[ind]))
            d_wave+=AMPA_d_amp[ind]*(np.exp(-t/AMPA_d_dlt[ind])-np.exp(-t/AMPA_d_rho))
            d_comp=(AMPA_d_amp[ind]*(np.exp(-t/AMPA_d_dlt[ind])-np.exp(-t/AMPA_d_rho)))
            #plt.plot(t, d_comp, label='d%s'%ind)
        #plt.legend()
        #plt.show()
        s_wave=np.zeros(len(t))
        for ind,_ in enumerate(AMPA_s_amp):
            #s_wave+=AMPA_s_amp[ind]*(np.exp(-t/AMPA_s_rho)-np.exp(-t/AMPA_s_dlt[ind]))
            s_wave+=AMPA_s_amp[ind]*(np.exp(-t/AMPA_s_dlt[ind])-np.exp(-t/AMPA_s_rho))
            #s_comp=(AMPA_s_amp[ind]*(np.exp(-t/AMPA_s_dlt[ind])-np.exp(-t/AMPA_s_rho)))
            #plt.plot(t, s_comp, label='s%s'%ind)
        #plt.legend()
        #plt.show()
        #sys.exit()

        #print('t fo AMPA', t)
        #plt.plot(t, s_wave, label="spillover")
        #plt.plot(t, d_wave, label="direct")
        #plt.title("AMPA wave for a spike")
        #plt.legend()
        #plt.show()
        #sys.exit()
        return self.p_AMPA_d(index_dend)*d_wave + self.p_AMPA_s(index_dend)*s_wave, r_AMPA_d*d_wave + r_AMPA_s*s_wave

    def block(self, STP=True):
        C1=2.07  *1e-03       #scaling term
        C2=0.015 *1e-03
        MG_out=1 *1e-03
        z, T, F_per_R, delta_bind, delta_perm= 2, 308.15, 96485.3329/8.31446261815324, 0.35, 0.53  #scale free
        theta=z*F_per_R/T
        Current_Mem_Pot=self.membrane_potential if STP else self.MEM_POT_NOSTP
        theta_V= theta*Current_Mem_Pot
        block = (
            C1*np.exp(delta_bind*theta_V)+C2*np.exp(-delta_perm*theta_V)
            )\
                /(
                    C1*np.exp(delta_bind*theta_V)+C2*np.exp(-delta_perm*theta_V)\
                    +MG_out*np.exp(-delta_bind*theta_V)
                    )
        import math
        if math.isnan(block) or block>1 or block<0:
            raise Exception('math.isnan(block) or block>1 or block<0\n',
            'block val:', block, 'term: c1', C1*np.exp(delta_bind*theta), 'c2', C2*np.exp(-delta_perm*theta),
            'mg', MG_out*np.exp(-delta_bind*theta), 'theta', theta, 'Mem Voltage:',Current_Mem_Pot)
        return block

    def NMDA_conductance_evaluation(self, t, index_dend):
        NMDA_wave=np.zeros(len(t))
        for ind, _ in enumerate(NMDA_amp):
            #NMDA_wave+=NMDA_amp[ind]*abs(np.exp(-t/NMDA_rho)-np.exp(-t/NMDA_dlt[ind]))
            NMDA_wave+=NMDA_amp[ind]*(np.exp(-t/NMDA_dlt[ind])-np.exp(-t/NMDA_rho))
        return self.p_NMDA(index_dend)*NMDA_wave*self.block(STP=True), r_NMDA*NMDA_wave*self.block(STP=False)

    def GABA_conductance_evaluation(self, t, index_dend):
        GABA_wave=np.zeros(len(t))
        for ind, _ in enumerate(GABA_amp):
            GABA_wave+=GABA_amp[ind]*(np.exp(-t/GABA_dlt[ind])-np.exp(-t/GABA_rho))
        return self.p_GABA(index_dend)*GABA_wave, r_GABA*GABA_wave

    def rate_of_change_membrane_voltage(self, Current_Mem_Pot, current_conductance_AMPA, current_conductance_NMDA,\
                                            current_condunctance_GABA, STP=True):

        dV=-1/membrane_capacitance*(\
                leak_conductance*(Current_Mem_Pot-leak_reversal_potential)\
                #+GABA_conductance*(Current_Mem_Pot-GABA_reversal_potential)\
                +current_condunctance_GABA*(Current_Mem_Pot-GABA_reversal_potential)\
                +current_conductance_AMPA *(Current_Mem_Pot-AMPA_reversal_potential)\
                #block* moved to NMDA evaluation function
                +current_conductance_NMDA*(Current_Mem_Pot-NMDA_reversal_potential)\
                )

        if STP: self.membrane_potential+=dV*TIME_UNIT  # change per milli second
        else: self.MEM_POT_NOSTP+=dV*TIME_UNIT

from RateCodingModel import *
from matplotlib import pyplot as plt
time_step=1000
TIME_RANGE=np.linspace(0,TIME_UNIT*time_step, time_step)
NUM_MFs=4
spiking_rate=80

Spike_trains=[]
Spike_trains+=poisson_spike_generator(80, time_step, 1e-03)  # Firing rate spkies per sec,
Spike_trains+=poisson_spike_generator(80, time_step, 1e-03)
Spike_trains+=poisson_spike_generator(100, time_step, 1e-03)
Spike_trains+=poisson_spike_generator(10, time_step, 1e-03)

spike_pttn_per_bin=np.array(Spike_trains).reshape(NUM_MFs, time_step)

#print(np.shape(spike_pttn_per_bin))
#print(np.shape(spike_pttn_per_bin[0]))
#sys.exit()

'''Input spike plot'''
#raster_plot(spike_pttn_per_bin, [0])
#sys.exit()

SPK_GC=Spiking_cells(num_dend=NUM_MFs)
''' For recording the values of variables as np.zeros'''
AMPA_Conductance=np.zeros((NUM_MFs, time_step))
NMDA_Conductance=np.zeros((NUM_MFs, time_step))
GABA_Conductance=np.zeros((NUM_MFs, time_step))

AMPA_NO_STP=np.zeros((NUM_MFs, time_step))
NMDA_NO_STP=np.zeros((NUM_MFs, time_step))
mem_voltage=np.zeros(time_step)
REFRACTORY_COUNT=0

mem_voltage2=np.zeros(time_step)
REFRACTORY_COUNT2=0

p_ofAMPA_d=np.zeros((NUM_MFs,time_step))
p_ofAMPA_s=np.zeros((NUM_MFs,time_step))
p_ofNMDA  =np.zeros((NUM_MFs,time_step))
p_ofGABA  =np.zeros((NUM_MFs,time_step))
''' Simulation loop start '''
for t, time in enumerate(TIME_RANGE):
    mem_voltage[t]=SPK_GC.membrane_potential

    for dend in range(SPK_GC.num_dend):
        p_ofAMPA_d[dend][t]=SPK_GC.p_AMPA_d(dend)
        p_ofAMPA_s[dend][t]=SPK_GC.p_AMPA_s(dend)
        p_ofNMDA[dend][t]=  SPK_GC.p_NMDA(dend)
        p_ofGABA[dend][t]=  SPK_GC.p_GABA(dend)

    if not REFRACTORY_COUNT==0:  #Refractory Period
        SPK_GC.membrane_potential=reset_membrane_potential
        REFRACTORY_COUNT-=1
    else:
        for dend in range(SPK_GC.num_dend):
            if spike_pttn_per_bin[dend][t]==1: #if there's a spike input for a dendrite
                AMPA_waveform_of_a_spike, _ =SPK_GC.AMPA_conductance_evaluation(np.arange(time_step-t), dend)
                NMDA_waveform_of_a_spike, _ =SPK_GC.NMDA_conductance_evaluation(np.arange(time_step-t), dend)
                GABA_waveform_of_a_spike, _ =SPK_GC.GABA_conductance_evaluation(np.arange(time_step-t), dend)
                AMPA_Conductance[dend][t:]+=AMPA_waveform_of_a_spike
                NMDA_Conductance[dend][t:]+=NMDA_waveform_of_a_spike
                GABA_Conductance[dend][t:]+=GABA_waveform_of_a_spike

            SPK_GC.rate_of_change_membrane_voltage(SPK_GC.membrane_potential,
                                                 AMPA_Conductance[dend][t], NMDA_Conductance[dend][t],\
                                                    GABA_Conductance[dend][t], STP=True)

        if SPK_GC.membrane_potential>=activation_threshold:
            print('SPIKE at time', t)
            REFRACTORY_COUNT=refractory_interval
            SPK_GC.membrane_potential=0          # 0mv Represent spikes (choosed as a arbitrary high value)
        elif SPK_GC.membrane_potential<=leak_reversal_potential:
            print('Limiting minimum potential by leak revers. Pot. at time', t)
            SPK_GC.membrane_potential=leak_reversal_potential

    for dend in range(SPK_GC.num_dend):
        SPK_GC.PLASTICITY(dend, spike_pttn_per_bin[dend][t]) #Short term plasticity

    # No STP simulation for Comparison-------------------------------
    '''
    mem_voltage2[t]=SPK_GC.MEM_POT_NOSTP

    if not REFRACTORY_COUNT2==0:  #Refractory Period
        SPK_GC.MEM_POT_NOSTP=reset_membrane_potential
        REFRACTORY_COUNT2-=1
    else:
        for dend in range(SPK_GC.num_dend):
            if spike_pttn_per_bin[dend][t]==1:
                _, AMPA_WAVE_No_STP=SPK_GC.AMPA_conductance_evaluation(np.arange(time_step-t), dend)
                _, NMDA_WAVE_No_STP=SPK_GC.NMDA_conductance_evaluation(np.arange(time_step-t), dend)
                AMPA_NO_STP[dend][t:]+=AMPA_WAVE_No_STP
                NMDA_NO_STP[dend][t:]+=NMDA_WAVE_No_STP

            SPK_GC.rate_of_change_membrane_voltage(SPK_GC.MEM_POT_NOSTP,
                                                AMPA_NO_STP[dend][t],NMDA_NO_STP[dend][t], STP=False)

        if SPK_GC.MEM_POT_NOSTP>=activation_threshold:
            print('SPIKE at time', t, 'For No STP')
            REFRACTORY_COUNT2=refractory_interval
            SPK_GC.MEM_POT_NOSTP=0   '''

#plt.plot(np.arange(time_step), conductance)
#plt.title('Conductance_ No_NMDA, NO STP')
#plt.show()

#raster_plot(spike_pttn_per_bin, [0])

Plotting_Scale=1 # of conducntance

print('Total spikes arrived:', np.sum(spike_pttn_per_bin), 'shape:', np.shape(spike_pttn_per_bin) )

fig, axs = plt.subplots(2+NUM_MFs, gridspec_kw={'height_ratios': [1]*(NUM_MFs)+[0.1]+[2]})

for ind_mf in range(NUM_MFs):
    axs[ind_mf].plot(np.arange(time_step), AMPA_Conductance[ind_mf]*Plotting_Scale, 'red', label="AMPA")
    axs[ind_mf].plot(np.arange(time_step), NMDA_Conductance[ind_mf]*Plotting_Scale, 'purple', label="NMDA")
    axs[ind_mf].plot(np.arange(time_step), GABA_Conductance[ind_mf]*Plotting_Scale, 'g', label="GABA")
    axs[ind_mf].axhline(y=630*1e-12*PARAM_SCALING, color='c', linestyle=':', label='Peak amplitude') #of both condunctances
    #axs[ind_mf].plot(np.arange(time_step), AMPA_NO_STP[ind_mf]*Plotting_Scale, label="AMPA_NO")
    #axs[ind_mf].plot(np.arange(time_step), NMDA_NO_STP[ind_mf]*Plotting_Scale, label="NMDA_NO")
    if ind_mf==3: axs[ind_mf].legend()

axs[0].set_title('Conductances')

#x,y = np.argwhere(spike_pttn_per_bin == 1).T
#axs[NUM_MFs].scatter(y, x, marker='|', label='SPIKE arrivals')
#axs[NUM_MFs].legend()

axs[NUM_MFs+1].plot(np.arange(time_step), mem_voltage, label='Mem. Pot.')
#axs[NUM_MFs].plot(np.arange(time_step), mem_voltage2, label='NO STP')
axs[NUM_MFs+1].axhline(y=reset_membrane_potential, color='r', linestyle=':', label='Reset_pot')
axs[NUM_MFs+1].axhline(y=activation_threshold, color='b', linestyle=':', label='Thrshld')
axs[NUM_MFs+1].set_ylabel("""Mmbrn Ptnt'l level""")
axs[NUM_MFs+1].legend()

plt.show()

#sys.exit()
# For P development on condunctance terms of STP
fig, axs = plt.subplots(NUM_MFs, gridspec_kw={'height_ratios': [1]*(NUM_MFs)})
for ind_mf in range(NUM_MFs):
    axs[ind_mf].plot(np.arange(time_step), p_ofAMPA_d[ind_mf], color='r', label="AMPA_d")
    axs[ind_mf].plot(np.arange(time_step), p_ofAMPA_s[ind_mf], color='c', label="AMPA_s")
    axs[ind_mf].plot(np.arange(time_step), p_ofNMDA[ind_mf],   color='m', label="NMDA")
    axs[ind_mf].plot(np.arange(time_step), p_ofGABA[ind_mf],   color='g', label="GABA")
    axs[ind_mf].axhline(y=r_AMPA_d, color='r', linestyle=':', label='Init_AMPA_d')
    axs[ind_mf].axhline(y=r_AMPA_s, color='c', linestyle=':', label='Init_AMPA_s')
    axs[ind_mf].axhline(y=r_NMDA,   color='m', linestyle=':', label='Init_NMDA')
    axs[ind_mf].axhline(y=r_GABA,   color='g', linestyle=':', label='Init_GABA')
    if ind_mf==3: axs[ind_mf].legend()

axs[0].set_title('P development')
plt.show()
sys.exit()

Plotting_Scale=1 # of conducntance
plt.plot(np.arange(time_step), AMPA_Conductance*Plotting_Scale, 'red', label="AMPA")
plt.plot(np.arange(time_step), NMDA_Conductance*Plotting_Scale, 'purple', label="NMDA")
#plt.title('Conductance, NO STP, No Mg block')
plt.title('Conductance, NO STP')
#plt.legend(handles=[plt_AMPA,plt_NMDA])
plt.legend()
plt.show()

#plt.plot(np.arange(time_step), mem_voltage/UNIT_SCALE)
plt.plot(np.arange(time_step), mem_voltage)
plt.title('Membrane Voltage')
plt.show()


from scripts.generator import Generator
from scripts.generator import Converter

class ArrayRBCMessage(Generator):
    def __init__(self, params: dict):
        super().__init__(['class_name', 'inher', 'guard', 'header_name', 'ros_msg_type', 'ros_msg_name',
                          'array_ros_type', 'array_ros_type_lower', 'array_descriptor'])
        self.header_template_name = 'templates/header_array.tmpl'
        self.impl_template_name = 'templates/impl_array.tmpl'
        self.param_dict = params
        self.param_dict['array_ros_type_lower'] = Converter.to_lower(d.get('array_ros_type'))

    def create(self):
        print("Creating files for class {}".format(self.param_dict.get('class_name')))
        super().read(self.header_template_name, self.impl_template_name)
        super().write(self.param_dict)

if __name__ == '__main__':
    d = {}
    d['class_name'] = 'PolygonMessage'
    d['ros_msg_type'] = 'geometry_msgs'
    d['ros_msg_name'] = 'Polygon'
    d['array_ros_type'] = 'Point'
    d['inher'] = 'override'
    d['array_descriptor'] = 'points'

    message = ArrayRBCMessage(d)
    message.create()


from edificio import Edificio

class casa(Edificio):
    Edificio._pisos = 2

    def calcularImpuestos(self, areaPiso):
        Edificio.areaDePiso = areaPiso
        areaTotal = Edificio._calcularAreaTotal()

        if (areaTotal < 200):
            return 0
        else:
            return 0.0012 * areaTotal

Curso Python UNED/Unidad 5/Ejercicio ejemplo/poo/casa.py


"""Test WGC angular frame transformation calculation."""

from pytest import fixture, raises

from webgeocalc import FrameTransformation
from webgeocalc.errors import (CalculationIncompatibleAttr,
                               CalculationInvalidAttr)

@fixture
def kernels():
    """Cassini kernel set."""
    return 5

@fixture
def time():
    """Input time."""
    return '2012-10-19T08:24:00.000'

@fixture
def frame_1():
    """Input frame 1."""
    return 'IAU_SATURN'

@fixture
def frame_2():
    """Input frame 2."""
    return 'IAU_ENCELADUS'

@fixture
def corr():
    """Input aberration correction."""
    return 'NONE'

@fixture
def params(kernels, time, frame_1, frame_2, corr):
    """Input parameters from WGC API example."""
    return {
        'kernels': kernels,
        'times': time,
        'frame_1': frame_1,
        'frame_2': frame_2,
        'aberration_correction': corr,
    }

@fixture
def payload(kernels, time, frame_1, frame_2, corr):
    """Payload from WGC API example."""
    return {
        "kernels": [{
            "type": "KERNEL_SET",
            "id": kernels,
        }],
        "timeSystem": "UTC",
        "timeFormat": "CALENDAR",
        "times": [
            time,
        ],
        "calculationType": "FRAME_TRANSFORMATION",
        "frame1": frame_1,
        "frame2": frame_2,
        "aberrationCorrection": corr,
        "timeLocation": "FRAME1",
        "orientationRepresentation": "EULER_ANGLES",
        "axis1": "X",
        "axis2": "Y",
        "axis3": "Z",
        "angularUnits": "deg",
        "angularVelocityRepresentation": "VECTOR_IN_FRAME1",
        "angularVelocityUnits": "deg/s"
    }

def test_frame_transformation_payload(params, payload):
    """Test angular frame transformation payload."""
    assert FrameTransformation(**params).payload == payload

def test_frame_transformation_attr_err(params):
    """Test errors when frame transformation is invalid."""
    del params['aberration_correction']
    with raises(CalculationInvalidAttr):
        # aberration_correctin can not be '+S'
        FrameTransformation(aberration_correction='CN+S', **params)

    with raises(CalculationInvalidAttr):
        FrameTransformation(time_location='WRONG', **params)

    with raises(CalculationInvalidAttr):
        FrameTransformation(orientation_representation='WRONG', **params)

    with raises(CalculationInvalidAttr):
        FrameTransformation(axis_1='WRONG', **params)

    with raises(CalculationInvalidAttr):
        FrameTransformation(angular_units='WRONG', **params)

    with raises(CalculationInvalidAttr):
        FrameTransformation(angular_velocity_representation='WRONG', **params)

    with raises(CalculationInvalidAttr):
        FrameTransformation(angular_velocity_units='WRONG', **params)

    with raises(CalculationIncompatibleAttr):
        FrameTransformation(angular_velocity_representation='EULER_ANGLE_DERIVATIVES',
                            angular_velocity_units='Unitary', **params)


#python gen.py --outdir example --ply_file ./example/3cl_pocket_8.0_res_1.5.ply
import os
import argparse
from glob import glob
from easydict import EasyDict
from rdkit import Chem
import torch
from copy import deepcopy
import shutil
import numpy as np
from tqdm.auto import tqdm
from utils.transforms import *
from utils.misc import load_config
from utils.reconstruct import *
from models.surfgen import SurfGen
from utils.chem import read_pkl, write_pkl
from utils.sample import get_init, get_next, logp_to_rank_prob, pdb_to_pocket_data
from utils.sample import STATUS_FINISHED, STATUS_RUNNING
import os.path as osp
from rdkit import RDLogger
lg = RDLogger.logger()
lg.setLevel(RDLogger.CRITICAL)
import warnings

import pickle
def write_pkl(list,file):
    with open(file,'wb') as f:
        pickle.dump(list,f)
        print('pkl file saved at {}'.format(file))
def read_pkl(file):
    with open(file,'rb') as f:
        data = pickle.load(f)
    return data

parser = argparse.ArgumentParser()
parser.add_argument(
    '--config', type=str, default='./configs/sample.yml'
)

parser.add_argument(
    '--outdir', type=str, default='example',
)

parser.add_argument(
    '--device', type=str, default='cuda',
)

parser.add_argument(
    '--check_point',type=str,default='./ckpt/surfgen.pt',
    help='load the parameter'
)

parser.add_argument(
    '--ply_file', action='store',required=False,type=str,default='./example/3cl_pocket_8.0_res_1.5.ply',
    help='surface file specified for generation'
)

args = parser.parse_args()
config = load_config(args.config)

data = pdb_to_pocket_data(args.ply_file)

contrastive_sampler = ContrastiveSample()
ligand_featurizer = FeaturizeLigandAtom()
protein_featurizer = FeaturizeProteinAtom()
transform = Compose([
    RefineData(),
    LigandCountNeighbors(),
    Geodesic_builder(),
    ligand_featurizer,
    protein_featurizer
])

ckpt = torch.load(args.check_point, map_location=args.device)
model = SurfGen(
    ckpt['config'].model,
    num_classes = contrastive_sampler.num_elements,
    num_bond_types = 3,
    protein_atom_feature_dim = protein_featurizer.feature_dim,
    ligand_atom_feature_dim = ligand_featurizer.feature_dim,
).to(args.device)
model.load_state_dict(ckpt['model'])
print('Num of parameters is {0:.4}M'.format(np.sum([p.numel() for p in model.parameters()]) /100000 ))

mask = LigandMaskAll()
composer = AtomComposer(5, ligand_featurizer.feature_dim, ckpt['config'].model.encoder.knn)
masking = Compose([
    mask,
    composer
])
def transform_data(data, transform):
    assert data.protein_pos.size(0) > 0
    if transform is not None:
        data = transform(data)
    return data

data = transform(data)
data = transform_data(data, masking)
np.seterr(invalid='ignore')
pool = EasyDict({
    'queue': [],
    'failed': [],
    'finished': [],
    'duplicate': [],
    'smiles': set(),
})

data = transform_data(deepcopy(data), masking)
init_data_list = get_init(data.to(args.device),   # sample the initial atoms
        model = model,
        transform=composer,
        threshold=config.sample.threshold
)
pool.queue = init_data_list
#rint('Start to generate novel molecules with 3D conformation located in the protein pocket!')
#print('The protein pocket is {}, init length is {}'.format(data.protein_filename, len(init_data_list)))
global_step = 0
while len(pool.finished) < config.sample.num_samples:
    global_step += 1
    if global_step > config.sample.max_steps:
        break
    queue_size = len(pool.queue)
    # # sample candidate new mols from each parent mol
    queue_tmp = []
    for data in pool.queue:
        nexts = []
        data_next_list = get_next(
            data.to(args.device),
            model = model,
            transform = composer,
            threshold = config.sample.threshold
        )

        for data_next in data_next_list:
            if data_next.status == STATUS_FINISHED:
                try:
                    rdmol = reconstruct_from_generated_with_edges(data_next)
                    data_next.rdmol = rdmol
                    mol = Chem.MolFromSmiles(Chem.MolToSmiles(rdmol))
                    smiles = Chem.MolToSmiles(mol)
                    data_next.smiles = smiles
                    if smiles in pool.smiles:
                        pool.duplicate.append(data_next)
                    elif '.' in smiles:
                        pool.failed.append(data_next)
                    else:   # Pass checks
                        print('Success: %s' % smiles)
                        pool.finished.append(data_next)
                        pool.smiles.add(smiles)
                except MolReconsError:
                    pool.failed.append(data_next)
            elif data_next.status == STATUS_RUNNING:
                nexts.append(data_next)

        queue_tmp += nexts
    prob = logp_to_rank_prob([p.average_logp[2:] for p in queue_tmp],)  # (logp_focal, logpdf_pos), logp_element, logp_hasatom, logp_bond
    n_tmp = len(queue_tmp)
    if n_tmp == 0:
        print('{}th has filures!'.format(global_step))
        break
    else:
        next_idx = np.random.choice(np.arange(n_tmp), p=prob, size=min(config.sample.beam_size, n_tmp), replace=False)
    pool.queue = [queue_tmp[idx] for idx in next_idx]

# save the generation results
task_name = args.ply_file.split('/')[-1][:-4]
task_dir = osp.join(args.outdir,task_name)
os.makedirs(task_dir,exist_ok=True)
sdf_file = os.path.join(task_dir,f'{task_name}_gen.sdf')
writer = Chem.SDWriter(sdf_file)
for j in range(len(pool['finished'])):
    writer.write(pool['finished'][j].rdmol)
writer.close()

SDF_dir = osp.join(task_dir,'SDF')
os.makedirs(SDF_dir, exist_ok=True)
for j in range(len(pool['finished'])):
    writer = Chem.SDWriter(SDF_dir+f'/{j}.sdf')
    writer.write(pool['finished'][j].rdmol)
    writer.close()

shutil.copy(args.ply_file,task_dir)

print('sucessfully generate at',task_dir)


import os
from copy import deepcopy

import torch
import torch.optim as optim
from torch.autograd import Variable
from torchvision import models

import CNN_model
from _utils import prediction, mean_average_precision

def CalcSim(batch_label, train_label):
    S = (batch_label.mm(train_label.t()) > 0).type(torch.FloatTensor)
    return S

def CreateModel(model_name, bit, use_gpu):
    if model_name == 'vgg11':
        vgg11 = models.vgg11(pretrained=True)
        cnn_model = CNN_model.cnn_model(vgg11, model_name, bit)
    if model_name == 'alexnet':
        alexnet = models.alexnet(pretrained=True)
        cnn_model = CNN_model.cnn_model(alexnet, model_name, bit)
    if use_gpu:
        cnn_model = cnn_model.cuda()
    return cnn_model

def AdjustLearningRate(optimizer, epoch, learning_rate):
    lr = learning_rate * (0.1 ** (epoch // 50))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return optimizer

def Logtrick(x, use_gpu):
    if use_gpu:
        lt = torch.log(1 + torch.exp(-torch.abs(x))) + torch.max(x, Variable(torch.FloatTensor([0.]).cuda()))
    else:
        lt = torch.log(1 + torch.exp(-torch.abs(x))) + torch.max(x, Variable(torch.FloatTensor([0.])))
    return lt

def Totloss(U, B, Sim, lamda, num_train):
    theta = U.mm(U.t()) / 2
    t1 = (theta * theta).sum() / (num_train * num_train)
    l1 = (- theta * Sim + Logtrick(Variable(theta), False).data).sum()
    l2 = (U - B).pow(2).sum()
    l = l1 + lamda * l2
    return l, l1, l2, t1

def DPSH_algo(param, train_loader, test_loader, database_loader, logger, gpu_ind=0):
    # parameters setting
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_ind)

    learning_rate = param["learning_rate"]
    weight_decay = param["weight_decay"]
    lamda = param['lambda']
    use_gpu = torch.cuda.is_available()

    ### create model
    model = CreateModel(param['model'], param['bit'], use_gpu)
    optimizer = optim.SGD(
        model.parameters(),
        lr=learning_rate,
        weight_decay=weight_decay)

    ### training phase
    # parameters setting
    num_train = len(train_loader.dataset)
    B = torch.zeros(num_train, param['bit'])
    U = torch.zeros(num_train, param['bit'])
    train_labels = train_loader.dataset.get_onehot_targets()

    totloss_record = []
    totl1_record = []
    totl2_record = []
    t1_record = []

    Sim = CalcSim(train_labels, train_labels)

    best_map = 0.0
    best_epoch = 0
    best_checkpoint = None
    count = 0
    for epoch in range(param['epochs']):
        model.train()
        epoch_loss = 0.0
        ## training epoch
        for train_input, train_label, batch_ind in train_loader:
            if use_gpu:
                S = CalcSim(train_label, train_labels)
                train_input = Variable(train_input.cuda())
            else:
                S = CalcSim(train_label, train_labels)
                train_input = Variable(train_input)

            model.zero_grad()
            train_outputs = model(train_input)
            for i, ind in enumerate(batch_ind):
                U[ind, :] = train_outputs.data[i]
                B[ind, :] = torch.sign(train_outputs.data[i])

            Bbatch = torch.sign(train_outputs)
            if use_gpu:
                theta_x = train_outputs.mm(Variable(U.cuda()).t()) / 2
                logloss = (Variable(S.cuda()) * theta_x - Logtrick(theta_x, use_gpu)).sum() / (
                        num_train * len(train_label))
                regterm = (Bbatch - train_outputs).pow(2).sum() / (num_train * len(train_label))
            else:
                theta_x = train_outputs.mm(Variable(U).t()) / 2
                logloss = (Variable(S) * theta_x - Logtrick(theta_x, use_gpu)).sum() / (num_train * len(train_label))
                regterm = (Bbatch - train_outputs).pow(2).sum() / (num_train * len(train_label))

            loss = - logloss + lamda * regterm
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()

        logger.info('[Train Phase][Epoch: %3d/%3d][Loss: %3.5f]' % (epoch, param['epochs'] - 1, epoch_loss / num_train))
        optimizer = AdjustLearningRate(optimizer, epoch, learning_rate)

        l, l1, l2, t1 = Totloss(U, B, Sim, lamda, num_train)
        totloss_record.append(l)
        totl1_record.append(l1)
        totl2_record.append(l2)
        t1_record.append(t1)

        logger.info('[Total Loss: %10.5f][total L1: %10.5f][total L2: %10.5f][norm theta: %3.5f]' % (l, l1, l2, t1))

        ## testing during epoch
        if (epoch + 1) % 5 == 0 or (epoch + 1) == param['epochs']:
            qB, qL = prediction(model, test_loader)
            rB, rL = prediction(model, database_loader)
            map = mean_average_precision(qB, rB, qL, rL, param['topk'])
            logger.info(
                f"[Evaluation][dataset:{param['dataset']}][bits:{param['bit']}][epoch:{epoch}/{param['epochs'] - 1}][best-mAP@{param['topk']}:{best_map:.7f}][mAP@{param['topk']}:{map:.7f}][count:{0 if map > best_map else (count + 1)}]")

            if map > best_map:
                best_map = map
                best_epoch = epoch
                best_checkpoint = deepcopy(model.state_dict())
                count = 0
            else:
                count += 1
                if count == 10:
                    logger.info(
                        f"without improvement, will save & exit, best mAP: {best_map}, best epoch: {best_epoch}")
                    torch.save(best_checkpoint, f"{param['save_dir']}/e{best_epoch}_{best_map:.3f}.pkl")
                    break
    if count != 10:
        logger.info(f"reach epoch limit, will save & exit, best mAP: {best_map}, best epoch: {best_epoch}")
        torch.save(best_checkpoint, f"{param['save_dir']}/e{best_epoch}_{best_map:.3f}.pkl")

    return best_epoch, best_map