"""
Contains the core building blocks of the framework.
"""
import math
from copy import deepcopy

import pandas as pd
import numpy as np
import cython as cy

class Node(object):

    """
    The Node is the main building block in bt's tree structure design.
    Both StrategyBase and SecurityBase inherit Node. It contains the
    core functionality of a tree node.

    Args:
        * name (str): The Node name
        * parent (Node): The parent Node
        * children (dict, list): A collection of children. If dict,
            the format is {name: child}, if list then list of children.

    Attributes:
        * name (str): Node name
        * parent (Node): Node parent
        * root (Node): Root node of the tree (topmost node)
        * children (dict): Node's children
        * now (datetime): Used when backtesting to store current date
        * stale (bool): Flag used to determine if Node is stale and need
            updating
        * prices (TimeSeries): Prices of the Node. Prices for a security will
            be the security's price, for a strategy it will be an index that
            reflects the value of the strategy over time.
        * price (float): last price
        * value (float): last value
        * weight (float): weight in parent
        * full_name (str): Name including parents' names
        * members (list): Current Node + node's children

    """

    _price = cy.declare(cy.double)
    _value = cy.declare(cy.double)
    _weight = cy.declare(cy.double)
    _issec = cy.declare(cy.bint)
    _has_strat_children = cy.declare(cy.bint)

    def __init__(self, name, parent=None, children=None):

        self.name = name

        # strategy children helpers
        self._has_strat_children = False
        self._strat_children = []

        # if children is not None, we assume that we want to limit the
        # available children space to the provided list.
        if children is not None:
            if isinstance(children, list):
                # if all strings - just save as universe_filter
                if all(isinstance(x, str) for x in children):
                    self._universe_tickers = children
                    # empty dict - don't want to uselessly create
                    # tons of children when they might not be needed
                    children = {}
                else:
                    # this will be case if we pass in children
                    # (say a bunch of sub-strategies)
                    tmp = {}
                    ut = []
                    for c in children:
                        if type(c) == str:
                            tmp[c] = SecurityBase(c)
                            ut.append(c)
                        else:
                            # deepcopy object for possible later reuse
                            tmp[c.name] = deepcopy(c)

                            # if strategy, turn on flag and add name to list
                            # strategy children have special treatment
                            if isinstance(c, StrategyBase):
                                self._has_strat_children = True
                                self._strat_children.append(c.name)
                            # if not strategy, then we will want to add this to
                            # universe_tickers to filter on setup
                            else:
                                ut.append(c.name)

                    children = tmp
                    # we want to keep whole universe in this case
                    # so set to None
                    self._universe_tickers = ut

        if parent is None:
            self.parent = self
            self.root = self
        else:
            self.parent = parent
            self.root = parent.root
            parent._add_child(self)

        # default children
        if children is None:
            children = {}
            self._universe_tickers = None
        self.children = children

        self._childrenv = children.values()
        for c in self._childrenv:
            c.parent = self
            c.root = self.root

        # set default value for now
        self.now = 0
        # make sure root has stale flag
        # used to avoid unncessary update
        # sometimes we change values in the tree and we know that we will need
        # to update if another node tries to access a given value (say weight).
        # This avoid calling the update until it is actually needed.
        self.root.stale = False

        # helper vars
        self._price = 0
        self._value = 0
        self._weight = 0

        # is security flag - used to avoid updating 0 pos securities
        self._issec = False

    def __getitem__(self, key):
        return self.children[key]

    @property
    def prices(self):
        """
        A TimeSeries of the Node's price.
        """
        # can optimize depending on type -
        # securities don't need to check stale to
        # return latest prices, whereas strategies do...
        raise NotImplementedError()

    @property
    def price(self):
        """
        Current price of the Node
        """
        # can optimize depending on type -
        # securities don't need to check stale to
        # return latest prices, whereas strategies do...
        raise NotImplementedError()

    @property
    def value(self):
        """
        Current value of the Node
        """
        if self.root.stale:
            self.root.update(self.root.now, None)
        return self._value

    @property
    def weight(self):
        """
        Current weight of the Node (with respect to the parent).
        """
        if self.root.stale:
            self.root.update(self.root.now, None)
        return self._weight

    def setup(self, dates):
        """
        Setup method used to initialize a Node with a set of dates.
        """
        raise NotImplementedError()

    def _add_child(self, child):
        child.parent = self
        child.root = self.root
        if self.children is None:
            self.children = {child.name: child}
        else:
            self.children[child.name] = child

        self._childrenv = self.children.values()

    def update(self, date, data=None, inow=None):
        """
        Update Node with latest date, and optionally some data.
        """
        raise NotImplementedError()

    def adjust(self, amount, update=True, isflow=True):
        """
        Adjust Node value by amount.
        """
        raise NotImplementedError()

    def allocate(self, amount, update=True):
        """
        Allocate capital to Node.
        """
        raise NotImplementedError()

    @property
    def members(self):
        """
        Node members. Members include current node as well as Node's
        children.
        """
        res = [self]
        for c in self.children.values():
            res.extend(c.members)
        return res

    @property
    def full_name(self):
        if self.parent == self:
            return self.name
        else:
            return '%s>%s' % (self.parent.full_name, self.name)

class StrategyBase(Node):

    """
    Strategy Node. Used to define strategy logic within a tree.
    A Strategy's role is to allocate capital to it's children
    based on a function.

    Args:
        * name (str): Strategy name
        * children (dict, list): A collection of children. If dict,
            the format is {name: child}, if list then list of children.
            Children can be any type of Node.
        * parent (Node): The parent Node

    Attributes:
        * name (str): Strategy name
        * parent (Strategy): Strategy parent
        * root (Strategy): Root node of the tree (topmost node)
        * children (dict): Strategy's children
        * now (datetime): Used when backtesting to store current date
        * stale (bool): Flag used to determine if Strategy is stale and need
            updating
        * prices (TimeSeries): Prices of the Strategy - basically an index that
            reflects the value of the strategy over time.
        * price (float): last price
        * value (float): last value
        * weight (float): weight in parent
        * full_name (str): Name including parents' names
        * members (list): Current Strategy + strategy's children
        * commission_fn (fn(quantity, price)): A function used to determine the
            commission (transaction fee) amount. Could be used to model slippage
            (implementation shortfall). Note that often fees are symmetric for
            buy and sell and absolute value of quantity should be used for
            calculation.
        * capital (float): Capital amount in Strategy - cash
        * universe (DataFrame): Data universe available at the current time.
            Universe contains the data passed in when creating a Backtest. Use
            this data to determine strategy logic.

    """

    _capital = cy.declare(cy.double)
    _net_flows = cy.declare(cy.double)
    _last_value = cy.declare(cy.double)
    _last_price = cy.declare(cy.double)
    _last_fee = cy.declare(cy.double)
    _paper_trade = cy.declare(cy.bint)
    bankrupt = cy.declare(cy.bint)

    def __init__(self, name, children=None, parent=None):
        Node.__init__(self, name, children=children, parent=parent)
        self._capital = 0
        self._weight = 1
        self._value = 0
        self._price = 100

        # helper vars
        self._net_flows = 0
        self._last_value = 0
        self._last_price = 100
        self._last_fee = 0

        # default commission function
        self.commission_fn = self._dflt_comm_fn

        self._paper_trade = False
        self._positions = None
        self.bankrupt = False

    @property
    def price(self):
        """
        Current price.
        """
        if self.root.stale:
            self.root.update(self.now, None)
        return self._price

    @property
    def prices(self):
        """
        TimeSeries of prices.
        """
        if self.root.stale:
            self.root.update(self.now, None)
        return self._prices.ix[:self.now]

    @property
    def values(self):
        """
        TimeSeries of values.
        """
        if self.root.stale:
            self.root.update(self.now, None)
        return self._values.ix[:self.now]

    @property
    def capital(self):
        """
        Current capital - amount of unallocated capital left in strategy.
        """
        # no stale check needed
        return self._capital

    @property
    def cash(self):
        """
        TimeSeries of unallocated capital.
        """
        # no stale check needed
        return self._cash

    @property
    def fees(self):
        """
        TimeSeries of fees.
        """
        # no stale check needed
        return self._fees

    @property
    def universe(self):
        """
        Data universe available at the current time.
        Universe contains the data passed in when creating a Backtest.
        Use this data to determine strategy logic.
        """
        # avoid windowing every time
        # if calling and on same date return
        # cached value
        if self.now == self._last_chk:
            return self._funiverse
        else:
            self._last_chk = self.now
            self._funiverse = self._universe.ix[:self.now]
            return self._funiverse

    @property
    def positions(self):
        """
        TimeSeries of positions.
        """
        # if accessing and stale - update first
        if self.root.stale:
            self.root.update(self.root.now, None)

        if self._positions is not None:
            return self._positions
        else:
            vals = pd.DataFrame({x.name: x.positions for x in self.members
                                 if isinstance(x, SecurityBase)})
            self._positions = vals
            return vals

    def setup(self, universe):
        """
        Setup strategy with universe. This will speed up future calculations
        and updates.
        """
        # save full universe in case we need it
        self._original_data = universe

        # determine if needs paper trading
        # and setup if so
        if self is not self.parent:
            self._paper_trade = True
            self._paper_amount = 1000000

            paper = deepcopy(self)
            paper.parent = paper
            paper.root = paper
            paper._paper_trade = False
            paper.setup(self._original_data)
            paper.adjust(self._paper_amount)
            self._paper = paper

        # setup universe
        funiverse = universe

        if self._universe_tickers is not None:
            # if we have universe_tickers defined, limit universe to
            # those tickers
            valid_filter = list(set(universe.columns)
                                .intersection(self._universe_tickers))

            funiverse = universe[valid_filter].copy()

            # if we have strat children, we will need to create their columns
            # in the new universe
            if self._has_strat_children:
                for c in self._strat_children:
                    funiverse[c] = np.nan

            # must create to avoid pandas warning
            funiverse = pd.DataFrame(funiverse)

        self._universe = funiverse
        # holds filtered universe
        self._funiverse = funiverse
        self._last_chk = None

        # We're not bankrupt yet
        self.bankrupt = False

        # setup internal data
        self.data = pd.DataFrame(index=funiverse.index,
                                 columns=['price', 'value', 'cash', 'fees'],
                                 data=0.0)

        self._prices = self.data['price']
        self._values = self.data['value']
        self._cash = self.data['cash']
        self._fees = self.data['fees']

        # setup children as well - use original universe here - don't want to
        # pollute with potential strategy children in funiverse
        if self.children is not None:
            [c.setup(universe) for c in self._childrenv]

    @cy.locals(newpt=cy.bint, val=cy.double, ret=cy.double)
    def update(self, date, data=None, inow=None):
        """
        Update strategy. Updates prices, values, weight, etc.
        """
        # resolve stale state
        self.root.stale = False

        # update helpers on date change
        # also set newpt flag
        newpt = False
        if self.now == 0:
            newpt = True
        elif date != self.now:
            self._net_flows = 0
            self._last_price = self._price
            self._last_value = self._value
            self._last_fee = 0.0
            newpt = True

        # update now
        self.now = date
        if inow is None:
            if self.now == 0:
                inow = 0
            else:
                inow = self.data.index.get_loc(date)

        # update children if any and calculate value
        val = self._capital  # default if no children

        if self.children is not None:
            for c in self._childrenv:
                # avoid useless update call
                if c._issec and not c._needupdate:
                    continue
                c.update(date, data, inow)
                val += c.value

        if self.root == self:
            if (val < 0) and not self.bankrupt:
                # Declare a bankruptcy
                self.bankrupt = True
                self.flatten()

        # update data if this value is different or
        # if now has changed - avoid all this if not since it
        # won't change
        if newpt or self._value != val:
            self._value = val
            self._values.values[inow] = val

            try:
                ret = self._value / (self._last_value
                                     + self._net_flows) - 1
            except ZeroDivisionError:
                if self._value == 0:
                    ret = 0
                else:
                    raise ZeroDivisionError(
                        'Could not update %s. Last value '
                        'was %s and net flows were %s. Current'
                        'value is %s. Therefore, '
                        'we are dividing by zero to obtain the return '
                        'for the period.' % (self.name,
                                             self._last_value,
                                             self._net_flows,
                                             self._value))

            self._price = self._last_price * (1 + ret)
            self._prices.values[inow] = self._price

        # update children weights
        if self.children is not None:
            for c in self._childrenv:
                # avoid useless update call
                if c._issec and not c._needupdate:
                    continue
                try:
                    c._weight = c.value / val
                except ZeroDivisionError:
                    c._weight = 0.0

        # if we have strategy children, we will need to update them in universe
        if self._has_strat_children:
            for c in self._strat_children:
                # TODO: optimize ".loc" here as well
                self._universe.loc[date, c] = self.children[c].price

        # Cash should track the unallocated capital at the end of the day, so
        # we should update it every time we call "update".
        # Same for fees
        self._cash.values[inow] = self._capital
        self._fees.values[inow] = self._last_fee

        # update paper trade if necessary
        if newpt and self._paper_trade:
            self._paper.update(date)
            self._paper.run()
            self._paper.update(date)
            # update price
            self._price = self._paper.price
            self._prices.values[inow] = self._price

    @cy.locals(amount=cy.double, update=cy.bint, flow=cy.bint, fees=cy.double)
    def adjust(self, amount, update=True, flow=True, fee=0.0):
        """
        Adjust capital - used to inject capital to a Strategy. This injection
        of capital will have no effect on the children.

        Args:
            * amount (float): Amount to adjust by.
            * update (bool): Force update?
            * flow (bool): Is this adjustment a flow? Basically a flow will
                have an impact on the price index. Examples of flows are
                commissions.

        """
        # adjust capital
        self._capital += amount
        self._last_fee += fee

        # if flow - increment net_flows - this will not affect
        # performance. Commissions and other fees are not flows since
        # they have a performance impact
        if flow:
            self._net_flows += amount

        if update:
            # indicates that data is now stale and must
            # be updated before access
            self.root.stale = True

    @cy.locals(amount=cy.double, update=cy.bint)
    def allocate(self, amount, child=None, update=True):
        """
        Allocate capital to Strategy. By default, capital is allocated
        recursively down the children, proportionally to the children's
        weights.  If a child is specified, capital will be allocated
        to that specific child.

        Allocation also have a side-effect. They will deduct the same amount
        from the parent's "account" to offset the allocation. If there is
        remaining capital after allocation, it will remain in Strategy.

        Args:
            * amount (float): Amount to allocate.
            * child (str): If specified, allocation will be directed to child
                only. Specified by name.
            * update (bool): Force update.

        """
        # allocate to child
        if child is not None:
            if child not in self.children:
                c = SecurityBase(child)
                c.setup(self._universe)
                # update to bring up to speed
                c.update(self.now)
                # add child to tree
                self._add_child(c)

            # allocate to child
            self.children[child].allocate(amount)
        # allocate to self
        else:
            # adjust parent's capital
            # no need to update now - avoids repetition
            if self.parent == self:
                self.parent.adjust(-amount, update=False, flow=True)
            else:
                # do NOT set as flow - parent will be another strategy
                # and therefore should not incur flow
                self.parent.adjust(-amount, update=False, flow=False)

            # adjust self's capital
            self.adjust(amount, update=False, flow=True)

            # push allocation down to children if any
            # use _weight to avoid triggering an update
            if self.children is not None:
                [c.allocate(amount * c._weight, update=False)
                 for c in self._childrenv]

            # mark as stale if update requested
            if update:
                self.root.stale = True

    @cy.locals(delta=cy.double, weight=cy.double, base=cy.double)
    def rebalance(self, weight, child, base=np.nan, update=True):
        """
        Rebalance a child to a given weight.

        This is a helper method to simplify code logic. This method is used
        when we want to se the weight of a particular child to a set amount.
        It is similar to allocate, but it calculates the appropriate allocation
        based on the current weight.

        Args:
            * weight (float): The target weight. Usually between -1.0 and 1.0.
            * child (str): child to allocate to - specified by name.
            * base (float): If specified, this is the base amount all weight
                delta calculations will be based off of. This is useful when we
                determine a set of weights and want to rebalance each child
                given these new weights. However, as we iterate through each
                child and call this method, the base (which is by default the
                current value) will change. Therefore, we can set this base to
                the original value before the iteration to ensure the proper
                allocations are made.
            * update (bool): Force update?

        """
        # if weight is 0 - we want to close child
        if weight == 0:
            if child in self.children:
                return self.close(child)
            else:
                return

        # if no base specified use self's value
        if np.isnan(base):
            base = self.value

        # else make sure we have child
        if child not in self.children:
            c = SecurityBase(child)
            c.setup(self._universe)
            # update child to bring up to speed
            c.update(self.now)
            self._add_child(c)

        # allocate to child
        # figure out weight delta
        c = self.children[child]
        delta = weight - c.weight
        c.allocate(delta * base)

    def close(self, child):
        """
        Close a child position - alias for rebalance(0, child). This will also
        flatten (close out all) the child's children.

        Args:
            * child (str): Child, specified by name.
        """
        c = self.children[child]
        # flatten if children not None
        if c.children is not None and len(c.children) != 0:
            c.flatten()
        c.allocate(-c.value)

    def flatten(self):
        """
        Close all child positions.
        """
        # go right to base alloc
        [c.allocate(-c.value) for c in self._childrenv if c.value != 0]

    def run(self):
        """
        This is the main logic method. Override this method to provide some
        algorithm to execute on each date change. This method is called by
        backtester.
        """
        pass

    def set_commissions(self, fn):
        """
        Set commission (transaction fee) function.

        Args:
            fn (fn(quantity, price)): Function used to determine commission
            amount.

        """
        self.commission_fn = fn

        for c in self._childrenv:
            if isinstance(c, StrategyBase):
                c.set_commissions(fn)

    @cy.locals(q=cy.double, p=cy.double)
    def _dflt_comm_fn(self, q, p):
        return max(1, abs(q) * 0.01)

class SecurityBase(Node):

    """
    Security Node. Used to define a security within a tree.
    A Security's has no children. It simply models an asset that can be bought
    or sold.

    Args:
        * name (str): Security name
        * multiplier (float): security multiplier - typically used for
            derivatives.

    Attributes:
        * name (str): Security name
        * parent (Security): Security parent
        * root (Security): Root node of the tree (topmost node)
        * now (datetime): Used when backtesting to store current date
        * stale (bool): Flag used to determine if Security is stale and need
            updating
        * prices (TimeSeries): Security prices.
        * price (float): last price
        * value (float): last value - basically position * price * multiplier
        * weight (float): weight in parent
        * full_name (str): Name including parents' names
        * members (list): Current Security + strategy's children
        * position (float): Current position (quantity).

    """

    _last_pos = cy.declare(cy.double)
    _position = cy.declare(cy.double)
    multiplier = cy.declare(cy.double)
    _prices_set = cy.declare(cy.bint)
    _needupdate = cy.declare(cy.bint)

    @cy.locals(multiplier=cy.double)
    def __init__(self, name, multiplier=1):
        Node.__init__(self, name, parent=None, children=None)
        self._value = 0
        self._price = 0
        self._weight = 0
        self._position = 0
        self.multiplier = multiplier

        # opt
        self._last_pos = 0
        self._issec = True
        self._needupdate = True

    @property
    def price(self):
        """
        Current price.
        """
        # if accessing and stale - update first
        if self._needupdate or self.now != self.parent.now:
            self.update(self.root.now)
        return self._price

    @property
    def prices(self):
        """
        TimeSeries of prices.
        """
        # if accessing and stale - update first
        if self._needupdate or self.now != self.parent.now:
            self.update(self.root.now)
        return self._prices.ix[:self.now]

    @property
    def values(self):
        """
        TimeSeries of values.
        """
        # if accessing and stale - update first
        if self._needupdate or self.now != self.parent.now:
            self.update(self.root.now)
        if self.root.stale:
            self.root.update(self.root.now, None)
        return self._values.ix[:self.now]

    @property
    def position(self):
        """
        Current position
        """
        # no stale check needed
        return self._position

    @property
    def positions(self):
        """
        TimeSeries of positions.
        """
        # if accessing and stale - update first
        if self._needupdate:
            self.update(self.root.now)
        if self.root.stale:
            self.root.update(self.root.now, None)
        return self._positions.ix[:self.now]

    def setup(self, universe):
        """
        Setup Security with universe. Speeds up future runs.

        Args:
            * universe (DataFrame): DataFrame of prices with security's name as
                one of the columns.

        """
        # if we already have all the prices, we will store them to speed up
        # future udpates
        try:
            prices = universe[self.name]
        except KeyError:
            prices = None

        # setup internal data
        if prices is not None:
            self._prices = prices
            self.data = pd.DataFrame(index=universe.index,
                                     columns=['value', 'position'],
                                     data=0.0)
            self._prices_set = True
        else:
            self.data = pd.DataFrame(index=universe.index,
                                     columns=['price', 'value', 'position'])
            self._prices = self.data['price']
            self._prices_set = False

        self._values = self.data['value']
        self._positions = self.data['position']

    @cy.locals(prc=cy.double)
    def update(self, date, data=None, inow=None):
        """
        Update security with a given date and optionally, some data.
        This will update price, value, weight, etc.
        """
        # filter for internal calls when position has not changed - nothing to
        # do. Internal calls (stale root calls) have None data. Also want to
        # make sure date has not changed, because then we do indeed want to
        # update.
        if date == self.now and self._last_pos == self._position:
            return

        if inow is None:
            if date == 0:
                inow = 0
            else:
                inow = self.data.index.get_loc(date)

        # date change - update price
        if date != self.now:
            # update now
            self.now = date

            if self._prices_set:
                self._price = self._prices.values[inow]
            # traditional data update
            elif data is not None:
                prc = data[self.name]
                self._price = prc
                self._prices.values[inow] = prc

        self._positions.values[inow] = self._position
        self._last_pos = self._position

        self._value = self._position * self._price * self.multiplier
        self._values.values[inow] = self._value

        if self._weight == 0 and self._position == 0:
            self._needupdate = False

    @cy.locals(amount=cy.double, update=cy.bint, q=cy.double, outlay=cy.double)
    def allocate(self, amount, update=True):
        """
        This allocates capital to the Security. This is the method used to
        buy/sell the security.

        A given amount of shares will be determined on the current price, a
        commisison will be calculated based on the parent's commission fn, and
        any remaining capital will be passed back up  to parent as an
        adjustment.

        Args:
            * amount (float): Amount of adjustment.
            * update (bool): Force update?

        """

        # will need to update if this has been idle for a while...
        # update if needupdate or if now is stale
        # fetch parent's now since our now is stale
        if self._needupdate or self.now != self.parent.now:
            self.update(self.parent.now)

        # ignore 0 alloc
        # Note that if the price of security has dropped to zero, then it should
        # never be selected by SelectAll, SelectN etc. I.e. we should not open
        # the position at zero price. At the same time, we are able to close
        # it at zero price, because at that point amount=0.
        # Note also that we don't erase the position in an asset which price has
        # dropped to zero (though the weight will indeed be = 0)
        if amount == 0:
            return

        if self.parent is self or self.parent is None:
            raise Exception(
                'Cannot allocate capital to a parentless security')

        if self._price == 0 or np.isnan(self._price):
            raise Exception(
                'Cannot allocate capital to '
                '%s because price is 0 or nan as of %s'
                % (self.name, self.parent.now))

        # buy/sell
        # determine quantity - must also factor in commission
        # closing out?
        if amount == -self._value:
            q = -self._position
        else:
            if (self._position > 0) or ((self._position == 0) and (amount > 0)):
                # if we're going long or changing long position
                q = math.floor(amount / (self._price * self.multiplier))
            else:
                # if we're going short or changing short position
                q = math.ceil(amount / (self._price * self.multiplier))

        # if q is 0 nothing to do
        if q == 0 or np.isnan(q):
            return

        # this security will need an update, even if pos is 0 (for example if
        # we close the positions, value and pos is 0, but still need to do that
        # last update)
        self._needupdate = True

        # adjust position & value
        self._position += q

        # calculate proper adjustment for parent
        # parent passed down amount so we want to pass
        # -outlay back up to parent to adjust for capital
        # used
        outlay, fee = self.outlay(q)

        # call parent
        self.parent.adjust(-outlay, update=update, flow=False, fee=fee)

    @cy.locals(q=cy.double, p=cy.double)
    def commission(self, q, p):
        """
        Calculates the commission (transaction fee) based on quantity and price.
        Uses the parent's commission_fn.

        Args:
            * q (float): quantity
            * p (float): price

        """
        return self.parent.commission_fn(q, p)

    @cy.locals(q=cy.double)
    def outlay(self, q):
        """
        Determines the complete cash outlay (including commission) necessary
        given a quantity q.
        Second returning parameter is a commission itself.

        Args:
            * q (float): quantity

        """
        fee = self.commission(q, self._price * self.multiplier)
        full_outlay = q * self._price * self.multiplier + fee
        return full_outlay, fee

    def run(self):
        """
        Does nothing - securities have nothing to do on run.
        """
        pass

class Algo(object):

    """
    Algos are used to modularize strategy logic so that strategy logic becomes
    modular, composable, more testable and less error prone. Basically, the
    Algo should follow the unix philosophy - do one thing well.

    In practice, algos are simply a function that receives one argument, the
    Strategy (refered to as target) and are expected to return a bool.

    When some state preservation is necessary between calls, the Algo
    object can be used (this object). The __call___ method should be
    implemented and logic defined therein to mimic a function call. A
    simple function may also be used if no state preservation is neceesary.

    Args:
        * name (str): Algo name

    """

    def __init__(self, name=None):
        self._name = name

    @property
    def name(self):
        """
        Algo name.
        """
        if self._name is None:
            self._name = self.__class__.__name__
        return self._name

    def __call__(self, target):
        raise NotImplementedError("%s not implemented!" % self.name)

class AlgoStack(Algo):

    """
    An AlgoStack derives from Algo runs multiple Algos until a
    failure is encountered.

    The purpose of an AlgoStack is to group a logic set of Algos together. Each
    Algo in the stack is run. Execution stops if one Algo returns False.

    Args:
        * algos (list): List of algos.

    """

    def __init__(self, *algos):
        super(AlgoStack, self).__init__()
        self.algos = algos
        self.check_run_always = any(hasattr(x, 'run_always')
                                    for x in self.algos)

    def __call__(self, target):
        # normal runing mode
        if not self.check_run_always:
            for algo in self.algos:
                if not algo(target):
                    return False
            return True
        # run mode when at least one algo has a run_always attribute
        else:
            # store result in res
            # allows continuation to check for and run
            # algos that have run_always set to True
            res = True
            for algo in self.algos:
                if res:
                    res = algo(target)
                elif hasattr(algo, 'run_always'):
                    if algo.run_always:
                        algo(target)
            return res

class Strategy(StrategyBase):

    """
    Strategy expands on the StrategyBase and incorporates Algos.

    Basically, a Strategy is built by passing in a set of algos. These algos
    will be placed in an Algo stack and the run function will call the stack.

    Furthermore, two class attributes are created to pass data between algos.
    perm for permanent data, temp for temporary data.

    Args:
        * name (str): Strategy name
        * algos (list): List of Algos to be passed into an AlgoStack
        * children (dict, list): Children - useful when you want to create
            strategies of strategies

    Attributes:
        * stack (AlgoStack): The stack
        * temp (dict): A dict containing temporary data - cleared on each call
            to run. This can be used to pass info to other algos.
        * perm (dict): Permanent data used to pass info from one algo to
            another. Not cleared on each pass.

    """

    def __init__(self, name, algos=[], children=None):
        super(Strategy, self).__init__(name, children=children)
        self.stack = AlgoStack(*algos)
        self.temp = {}
        self.perm = {}

    def run(self):
        # clear out temp data
        self.temp = {}

        # run algo stack
        self.stack(self)

        # run children
        for c in self.children.values():
            c.run()


"""
Shogun demo, based on PyQT Demo by Eli Bendersky

Christian Widmer
Soeren Sonnenburg
License: GPLv3
"""
import numpy
import sys, os, csv
from PyQt4.QtCore import *
from PyQt4.QtGui import *

import matplotlib
from matplotlib.colorbar import make_axes, Colorbar
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure

from shogun import *
import util

class Form(QMainWindow):
    def __init__(self, parent=None):
        super(Form, self).__init__(parent)
        self.setWindowTitle('SHOGUN interactive demo')

        self.data = DataHolder()
        self.series_list_model = QStandardItemModel()

        self.create_menu()
        self.create_main_frame()
        self.create_status_bar()

        self.on_show()

    def load_file(self, filename=None):
        filename = QFileDialog.getOpenFileName(self,
            'Open a data file', '.', 'CSV files (*.csv);;All Files (*.*)')

        if filename:
            self.data.load_from_file(filename)
            self.fill_series_list(self.data.series_names())
            self.status_text.setText("Loaded " + filename)

    def on_show(self):
        self.axes.clear()
        self.axes.grid(True)
        self.axes.plot(self.data.x1_pos, self.data.x2_pos, 'ro')
        self.axes.plot(self.data.x1_neg, self.data.x2_neg, 'bo')
        self.axes.set_xlim((-5,5))
        self.axes.set_ylim((-5,5))
        self.canvas.draw()
        self.fill_series_list(self.data.get_stats())

    def on_about(self):
        msg = __doc__
        QMessageBox.about(self, "About the demo", msg.strip())

    def fill_series_list(self, names):
        self.series_list_model.clear()

        for name in names:
            item = QStandardItem(name)
            item.setCheckState(Qt.Unchecked)
            item.setCheckable(False)
            self.series_list_model.appendRow(item)

    def onclick(self, event):
        print 'button=%d, x=%d, y=%d, xdata=%f, ydata=%f'%(event.button, event.x, event.y, event.xdata, event.ydata)
        if event.button==1:
            label = 1.0
        else:
            label = -1.0
        self.data.add_example(event.xdata, event.ydata, label)
        self.on_show()

    def clear(self):
        self.data.clear()
        self.on_show()

    def enable_widgets(self):
        kernel_name = self.kernel_combo.currentText()
        if kernel_name == "LinearKernel":
            self.sigma.setDisabled(True)
            self.degree.setDisabled(True)
        elif kernel_name == "PolynomialKernel":
            self.sigma.setDisabled(True)
            self.degree.setEnabled(True)
        elif kernel_name == "GaussianKernel":
            self.sigma.setEnabled(True)
            self.degree.setDisabled(True)

    def train_svm(self):

        width = float(self.sigma.text())
        degree = int(self.degree.text())

        self.axes.clear()
        self.axes.grid(True)
        self.axes.plot(self.data.x1_pos, self.data.x2_pos, 'ro')
        self.axes.plot(self.data.x1_neg, self.data.x2_neg, 'bo')

        # train svm
        labels = self.data.get_labels()
        print type(labels)
        lab = BinaryLabels(labels)
        features = self.data.get_examples()
        train = RealFeatures(features)

        kernel_name = self.kernel_combo.currentText()
        print "current kernel is %s" % (kernel_name)

        if kernel_name == "LinearKernel":
            gk = LinearKernel(train, train)
            gk.set_normalizer(IdentityKernelNormalizer())
        elif kernel_name == "PolynomialKernel":
            gk = PolyKernel(train, train, degree, True)
            gk.set_normalizer(IdentityKernelNormalizer())
        elif kernel_name == "GaussianKernel":
            gk = GaussianKernel(train, train, width)

        cost = float(self.cost.text())

        print "cost", cost
        svm = LibSVM(cost, gk, lab)
        svm.train()
        svm.set_epsilon(1e-2)

        x, y, z = util.compute_output_plot_isolines(svm, gk, train)
        plt=self.axes.pcolor(x, y, z)
        CS=self.axes.contour(x, y, z, [-1,0,1], linewidths=1, colors='black', hold=True)
        #CS=self.axes.contour(x, y, z, linewidths=1, colors='black', hold=True)
        #CS=self.axes.contour(x, y, z, 5, linewidths=1, colors='black', hold=True)
        matplotlib.pyplot.clabel(CS, inline=1, fontsize=10)

        self.axes.set_xlim((-5,5))
        self.axes.set_ylim((-5,5))

        cmap = matplotlib.cm.jet
        norm = matplotlib.colors.Normalize(numpy.min(z), numpy.max(z))
        print CS.get_clim()
        if not self.cax:
            self.cax, kw = make_axes(self.axes)

# ColorbarBase derives from ScalarMappable and puts a colorbar
# in a specified axes, so it has everything needed for a
# standalone colorbar.  There are many more kwargs, but the
# following gives a basic continuous colorbar with ticks
# and labels.
        cb1 = matplotlib.colorbar.ColorbarBase(self.cax, cmap=cmap,
                                           norm=norm)
        self.canvas.draw()

    def create_main_frame(self):
        self.main_frame = QWidget()

        plot_frame = QWidget()

        self.dpi = 100
        self.fig = Figure((6.0, 6.0), dpi=self.dpi)
        self.canvas = FigureCanvas(self.fig)
        self.canvas.setParent(self.main_frame)

        cid = self.canvas.mpl_connect('button_press_event', self.onclick)
        self.axes = self.fig.add_subplot(111)
        self.cax = None
        #self.mpl_toolbar = NavigationToolbar(self.canvas, self.main_frame)

        log_label = QLabel("Number of examples:")
        self.series_list_view = QListView()
        self.series_list_view.setModel(self.series_list_model)

        cost_label = QLabel('C')
        #self.cost = QSpinBox()#QLineEdit()
        self.cost = QLineEdit()
        self.cost.setText("1.0")
        #self.cost.setMinimum(1)
        spin_label2 = QLabel('sigma')
        self.sigma = QLineEdit()
        self.sigma.setText("1.2")
        #self.sigma.setMinimum(1)
        self.degree = QLineEdit()
        self.degree.setText("2")
        #self.sigma.setMinimum(1)

        spins_hbox = QHBoxLayout()
        spins_hbox.addWidget(cost_label)
        spins_hbox.addWidget(self.cost)
        spins_hbox.addWidget(spin_label2)
        spins_hbox.addWidget(self.sigma)
        spins_hbox.addWidget(self.degree)
        spins_hbox.addStretch(1)

        self.legend_cb = QCheckBox("Show Support Vectors")
        self.legend_cb.setChecked(False)

        self.show_button = QPushButton("&Train SVM")
        self.connect(self.show_button, SIGNAL('clicked()'), self.train_svm)

        self.clear_button = QPushButton("&Clear")
        self.connect(self.clear_button, SIGNAL('clicked()'), self.clear)

        self.kernel_combo = QComboBox()
        self.kernel_combo.insertItem(-1, "GaussianKernel")
        self.kernel_combo.insertItem(-1, "PolynomialKernel")
        self.kernel_combo.insertItem(-1, "LinearKernel")
        self.kernel_combo.maximumSize = QSize(300, 50)
        self.connect(self.kernel_combo, SIGNAL("currentIndexChanged(QString)"), self.enable_widgets)

        left_vbox = QVBoxLayout()
        left_vbox.addWidget(self.canvas)
        #left_vbox.addWidget(self.mpl_toolbar)

        right0_vbox = QVBoxLayout()
        right0_vbox.addWidget(log_label)
        right0_vbox.addWidget(self.series_list_view)
        #right0_vbox.addWidget(self.legend_cb)
        right0_vbox.addStretch(1)

        right2_vbox = QVBoxLayout()
        right2_label = QLabel("Settings")
        right2_vbox.addWidget(right2_label)
        right2_vbox.addWidget(self.show_button)
        right2_vbox.addWidget(self.kernel_combo)
        right2_vbox.addLayout(spins_hbox)
        right2_clearlabel = QLabel("Remove Data")
        right2_vbox.addWidget(right2_clearlabel)

        right2_vbox.addWidget(self.clear_button)

        right2_vbox.addStretch(1)

        right_vbox = QHBoxLayout()
        right_vbox.addLayout(right0_vbox)
        right_vbox.addLayout(right2_vbox)

        hbox = QVBoxLayout()
        hbox.addLayout(left_vbox)
        hbox.addLayout(right_vbox)
        self.main_frame.setLayout(hbox)

        self.setCentralWidget(self.main_frame)
        self.enable_widgets()

    def create_status_bar(self):
        self.status_text = QLabel("")
        self.statusBar().addWidget(self.status_text, 1)

    def create_menu(self):
        self.file_menu = self.menuBar().addMenu("&File")

        load_action = self.create_action("&Load file",
            shortcut="Ctrl+L", slot=self.load_file, tip="Load a file")
        quit_action = self.create_action("&Quit", slot=self.close,
            shortcut="Ctrl+Q", tip="Close the application")

        self.add_actions(self.file_menu,
            (load_action, None, quit_action))

        self.help_menu = self.menuBar().addMenu("&Help")
        about_action = self.create_action("&About",
            shortcut='F1', slot=self.on_about,
            tip='About the demo')

        self.add_actions(self.help_menu, (about_action,))

    def add_actions(self, target, actions):
        for action in actions:
            if action is None:
                target.addSeparator()
            else:
                target.addAction(action)

    def create_action(  self, text, slot=None, shortcut=None,
                        icon=None, tip=None, checkable=False,
                        signal="triggered()"):
        action = QAction(text, self)
        if icon is not None:
            action.setIcon(QIcon(":/%s.png" % icon))
        if shortcut is not None:
            action.setShortcut(shortcut)
        if tip is not None:
            action.setToolTip(tip)
            action.setStatusTip(tip)
        if slot is not None:
            self.connect(action, SIGNAL(signal), slot)
        if checkable:
            action.setCheckable(True)
        return action

class DataHolder(object):
    """ Just a thin wrapper over a dictionary that holds integer
        data series. Each series has a name and a list of numbers
        as its data. The length of all series is assumed to be
        the same.

        The series can be read from a CSV file, where each line
        is a separate series. In each series, the first item in
        the line is the name, and the rest are data numbers.
    """
    def __init__(self, filename=None):
        self.clear()
        self.load_from_file(filename)

    def clear(self):
        self.x1_pos = []
        self.x2_pos = []
        self.x1_neg = []
        self.x2_neg = []

    def get_stats(self):

        num_neg = len(self.x1_neg)
        num_pos = len(self.x1_pos)

        str_neg = "num negative examples: %i" % num_neg
        str_pos = "num positive examples: %i" % num_pos

        return (str_neg, str_pos)

    def get_labels(self):
        return numpy.array([1]*len(self.x1_pos) + [-1]*len(self.x1_neg), dtype=numpy.float64)

    def get_examples(self):
        num_pos = len(self.x1_pos)
        num_neg = len(self.x1_neg)
        examples = numpy.zeros((2,num_pos+num_neg))

        for i in xrange(num_pos):
            examples[0,i] = self.x1_pos[i]
            examples[1,i] = self.x2_pos[i]

        for i in xrange(num_neg):
            examples[0,i+num_pos] = self.x1_neg[i]
            examples[1,i+num_pos] = self.x2_neg[i]

        return examples

    def add_example(self, x1, x2, label):
        if label==1:
            self.x1_pos.append(x1)
            self.x2_pos.append(x2)
        else:
            self.x1_neg.append(x1)
            self.x2_neg.append(x2)

    def load_from_file(self, filename=None):
        self.data = {}
        self.names = []

        if filename:
            for line in csv.reader(open(filename, 'rb')):
                self.names.append(line[0])
                self.data[line[0]] = map(int, line[1:])
                self.datalen = len(line[1:])

    def series_names(self):
        """ Names of the data series
        """
        return self.names

    def series_len(self):
        """ Length of a data series
        """
        return self.datalen

    def series_count(self):
        return len(self.data)

    def get_series_data(self, name):
        return self.data[name]

def main():
    app = QApplication(sys.argv)
    form = Form()
    form.show()
    app.exec_()

if __name__ == "__main__":
    main()

    #~ dh = DataHolder('qt_mpl_data.csv')
    #~ print dh.data
    #~ print dh.get_series_data('1991 Sales')
    #~ print dh.series_names()
    #~ print dh.series_count()


# Author: Travis Oliphant
from __future__ import division, print_function, absolute_import

__all__ = ['odeint']

from . import _odepack
from copy import copy
import warnings

class ODEintWarning(Warning):
    pass

_msgs = {2: "Integration successful.",
         1: "Nothing was done; the integration time was 0.",
         -1: "Excess work done on this call (perhaps wrong Dfun type).",
         -2: "Excess accuracy requested (tolerances too small).",
         -3: "Illegal input detected (internal error).",
         -4: "Repeated error test failures (internal error).",
         -5: "Repeated convergence failures (perhaps bad Jacobian or tolerances).",
         -6: "Error weight became zero during problem.",
         -7: "Internal workspace insufficient to finish (internal error)."
         }

def odeint(func, y0, t, args=(), Dfun=None, col_deriv=0, full_output=0,
           ml=None, mu=None, rtol=None, atol=None, tcrit=None, h0=0.0,
           hmax=0.0, hmin=0.0, ixpr=0, mxstep=0, mxhnil=0, mxordn=12,
           mxords=5, printmessg=0):
    """
    Integrate a system of ordinary differential equations.

    Solve a system of ordinary differential equations using lsoda from the
    FORTRAN library odepack.

    Solves the initial value problem for stiff or non-stiff systems
    of first order ode-s::

        dy/dt = func(y, t0, ...)

    where y can be a vector.

    *Note*: The first two arguments of ``func(y, t0, ...)`` are in the
    opposite order of the arguments in the system definition function used
    by the `scipy.integrate.ode` class.

    Parameters
    ----------
    func : callable(y, t0, ...)
        Computes the derivative of y at t0.
    y0 : array
        Initial condition on y (can be a vector).
    t : array
        A sequence of time points for which to solve for y.  The initial
        value point should be the first element of this sequence.
    args : tuple, optional
        Extra arguments to pass to function.
    Dfun : callable(y, t0, ...)
        Gradient (Jacobian) of `func`.
    col_deriv : bool, optional
        True if `Dfun` defines derivatives down columns (faster),
        otherwise `Dfun` should define derivatives across rows.
    full_output : bool, optional
        True if to return a dictionary of optional outputs as the second output
    printmessg : bool, optional
        Whether to print the convergence message

    Returns
    -------
    y : array, shape (len(t), len(y0))
        Array containing the value of y for each desired time in t,
        with the initial value `y0` in the first row.
    infodict : dict, only returned if full_output == True
        Dictionary containing additional output information

        =======  ============================================================
        key      meaning
        =======  ============================================================
        'hu'     vector of step sizes successfully used for each time step.
        'tcur'   vector with the value of t reached for each time step.
                 (will always be at least as large as the input times).
        'tolsf'  vector of tolerance scale factors, greater than 1.0,
                 computed when a request for too much accuracy was detected.
        'tsw'    value of t at the time of the last method switch
                 (given for each time step)
        'nst'    cumulative number of time steps
        'nfe'    cumulative number of function evaluations for each time step
        'nje'    cumulative number of jacobian evaluations for each time step
        'nqu'    a vector of method orders for each successful step.
        'imxer'  index of the component of largest magnitude in the
                 weighted local error vector (e / ewt) on an error return, -1
                 otherwise.
        'lenrw'  the length of the double work array required.
        'leniw'  the length of integer work array required.
        'mused'  a vector of method indicators for each successful time step:
                 1: adams (nonstiff), 2: bdf (stiff)
        =======  ============================================================

    Other Parameters
    ----------------
    ml, mu : int, optional
        If either of these are not None or non-negative, then the
        Jacobian is assumed to be banded.  These give the number of
        lower and upper non-zero diagonals in this banded matrix.
        For the banded case, `Dfun` should return a matrix whose
        rows contain the non-zero bands (starting with the lowest diagonal).
        Thus, the return matrix `jac` from `Dfun` should have shape
        ``(ml + mu + 1, len(y0))`` when ``ml >=0`` or ``mu >=0``.
        The data in `jac` must be stored such that ``jac[i - j + mu, j]``
        holds the derivative of the `i`th equation with respect to the `j`th
        state variable.  If `col_deriv` is True, the transpose of this
        `jac` must be returned.
    rtol, atol : float, optional
        The input parameters `rtol` and `atol` determine the error
        control performed by the solver.  The solver will control the
        vector, e, of estimated local errors in y, according to an
        inequality of the form ``max-norm of (e / ewt) <= 1``,
        where ewt is a vector of positive error weights computed as
        ``ewt = rtol * abs(y) + atol``.
        rtol and atol can be either vectors the same length as y or scalars.
        Defaults to 1.49012e-8.
    tcrit : ndarray, optional
        Vector of critical points (e.g. singularities) where integration
        care should be taken.
    h0 : float, (0: solver-determined), optional
        The step size to be attempted on the first step.
    hmax : float, (0: solver-determined), optional
        The maximum absolute step size allowed.
    hmin : float, (0: solver-determined), optional
        The minimum absolute step size allowed.
    ixpr : bool, optional
        Whether to generate extra printing at method switches.
    mxstep : int, (0: solver-determined), optional
        Maximum number of (internally defined) steps allowed for each
        integration point in t.
    mxhnil : int, (0: solver-determined), optional
        Maximum number of messages printed.
    mxordn : int, (0: solver-determined), optional
        Maximum order to be allowed for the non-stiff (Adams) method.
    mxords : int, (0: solver-determined), optional
        Maximum order to be allowed for the stiff (BDF) method.

    See Also
    --------
    ode : a more object-oriented integrator based on VODE.
    quad : for finding the area under a curve.

    Examples
    --------
    The second order differential equation for the angle `theta` of a
    pendulum acted on by gravity with friction can be written::

        theta''(t) + b*theta'(t) + c*sin(theta(t)) = 0

    where `b` and `c` are positive constants, and a prime (') denotes a
    derivative.  To solve this equation with `odeint`, we must first convert
    it to a system of first order equations.  By defining the angular
    velocity ``omega(t) = theta'(t)``, we obtain the system::

        theta'(t) = omega(t)
        omega'(t) = -b*omega(t) - c*sin(theta(t))

    Let `y` be the vector [`theta`, `omega`].  We implement this system
    in python as:

    >>> def pend(y, t, b, c):
    ...     theta, omega = y
    ...     dydt = [omega, -b*omega - c*np.sin(theta)]
    ...     return dydt
    ...

    We assume the constants are `b` = 0.25 and `c` = 5.0:

    >>> b = 0.25
    >>> c = 5.0

    For initial conditions, we assume the pendulum is nearly vertical
    with `theta(0)` = `pi` - 0.1, and it initially at rest, so
    `omega(0)` = 0.  Then the vector of initial conditions is

    >>> y0 = [np.pi - 0.1, 0.0]

    We generate a solution 101 evenly spaced samples in the interval
    0 <= `t` <= 10.  So our array of times is:

    >>> t = np.linspace(0, 10, 101)

    Call `odeint` to generate the solution.  To pass the parameters
    `b` and `c` to `pend`, we give them to `odeint` using the `args`
    argument.

    >>> from scipy.integrate import odeint
    >>> sol = odeint(pend, y0, t, args=(b, c))

    The solution is an array with shape (101, 2).  The first column
    is `theta(t)`, and the second is `omega(t)`.  The following code
    plots both components.

    >>> import matplotlib.pyplot as plt
    >>> plt.plot(t, sol[:, 0], 'b', label='theta(t)')
    >>> plt.plot(t, sol[:, 1], 'g', label='omega(t)')
    >>> plt.legend(loc='best')
    >>> plt.xlabel('t')
    >>> plt.grid()
    >>> plt.show()
    """

    if ml is None:
        ml = -1  # changed to zero inside function call
    if mu is None:
        mu = -1  # changed to zero inside function call
    t = copy(t)
    y0 = copy(y0)
    output = _odepack.odeint(func, y0, t, args, Dfun, col_deriv, ml, mu,
                             full_output, rtol, atol, tcrit, h0, hmax, hmin,
                             ixpr, mxstep, mxhnil, mxordn, mxords)
    if output[-1] < 0:
        warning_msg = _msgs[output[-1]] + " Run with full_output = 1 to get quantitative information."
        warnings.warn(warning_msg, ODEintWarning)
    elif printmessg:
        warning_msg = _msgs[output[-1]]
        warnings.warn(warning_msg, ODEintWarning)

    if full_output:
        output[1]['message'] = _msgs[output[-1]]

    output = output[:-1]
    if len(output) == 1:
        return output[0]
    else:
        return output


# coding: utf-8

# # Antibody Response Pulse
# https://github.com/blab/antibody-response-pulse
#
# ### B-cells evolution --- cross-reactive antibody response after influenza virus infection or vaccination
# ### Adaptive immune response for repeated infection

# In[3]:

'''
author: Alvason Zhenhua Li
date:   04/09/2015
'''
get_ipython().magic(u'matplotlib inline')

import numpy as np
import matplotlib.pyplot as plt
import os
from matplotlib.ticker import FuncFormatter

import alva_machinery_event_OAS_new as alva

AlvaFontSize = 23
AlvaFigSize = (15, 5)
numberingFig = 0

# plotting
dir_path = '/Users/al/Desktop/GitHub/antibody-response-pulse/bcell-array/figure'
file_name = 'Virus-Bcell-IgM-IgG'
figure_name = '-equation'
file_suffix = '.png'
save_figure = os.path.join(dir_path, file_name + figure_name + file_suffix)

numberingFig = numberingFig + 1
plt.figure(numberingFig, figsize=(12, 5))
plt.axis('off')
plt.title(r'$ Virus-Bcell-IgM-IgG \ equations \ (antibody-response \ for \ repeated-infection) $'
          , fontsize = AlvaFontSize)
plt.text(0, 7.0/9, r'$ \frac{\partial V_n(t)}{\partial t} =          +\mu_{v}V_{n}(t)(1 - \frac{V_n(t)}{V_{max}}) - \phi_{m} M_{n}(t) V_{n}(t) - \phi_{g} G_{n}(t) V_{n}(t) $'
         , fontsize = 1.2*AlvaFontSize)
plt.text(0, 5.0/9, r'$ \frac{\partial B_n(t)}{\partial t} =          +\mu_{b}V_{n}(t)(1 - \frac{V_n(t)}{V_{max}}) + (\beta_{m} + \beta_{g}) V_{n}(t) B_{n}(t) - \mu_{b} B_{n}(t)          + m_b V_{n}(t)\frac{B_{i-1}(t) - 2B_i(t) + B_{i+1}(t)}{(\Delta i)^2} $'
         , fontsize = 1.2*AlvaFontSize)
plt.text(0, 3.0/9,r'$ \frac{\partial M_n(t)}{\partial t} =          +\xi_{m} B_{n}(t) - \phi_{m} M_{n}(t) V_{n}(t) - \mu_{m} M_{n}(t) $'
         , fontsize = 1.2*AlvaFontSize)
plt.text(0, 1.0/9,r'$ \frac{\partial G_n(t)}{\partial t} =          +\xi_{g} B_{n}(t) - \phi_{g} G_{n}(t) V_{n}(t) - \mu_{g} G_{n}(t)          + m_a V_{n}(t)\frac{G_{i-1}(t) - 2G_i(t) + G_{i+1}(t)}{(\Delta i)^2} $'
         , fontsize = 1.2*AlvaFontSize)

plt.savefig(save_figure, dpi = 100)
plt.show()

# define the V-M-G partial differential equations
def dVdt_array(VBMGxt = [], *args):
    # naming
    V = VBMGxt[0]
    B = VBMGxt[1]
    M = VBMGxt[2]
    G = VBMGxt[3]
    x_totalPoint = VBMGxt.shape[1]
    # there are n dSdt
    dV_dt_array = np.zeros(x_totalPoint)
    # each dSdt with the same equation form
    dV_dt_array[:] = +inRateV*V[:]*(1 - V[:]/maxV) - killRateVm*M[:]*V[:] - killRateVg*G[:]*V[:]
    return(dV_dt_array)

def dBdt_array(VBMGxt = [], *args):
    # naming
    V = VBMGxt[0]
    B = VBMGxt[1]
    M = VBMGxt[2]
    G = VBMGxt[3]
    x_totalPoint = VBMGxt.shape[1]
    # there are n dSdt
    dB_dt_array = np.zeros(x_totalPoint)
    # each dSdt with the same equation form
    Bcopy = np.copy(B)
    centerX = Bcopy[:]
    leftX = np.roll(Bcopy[:], 1)
    rightX = np.roll(Bcopy[:], -1)
    leftX[0] = centerX[0]
    rightX[-1] = centerX[-1]
    dB_dt_array[:] = +inRateB*V[:]*(1 - V[:]/maxV) + (actRateBm + alva.event_active + alva.event_OAS_B)*V[:]*B[:] - outRateB*B[:]                      + mutatRateB*V[:]*(leftX[:] - 2*centerX[:] + rightX[:])/(dx**2)
    return(dB_dt_array)

def dMdt_array(VBMGxt = [], *args):
    # naming
    V = VBMGxt[0]
    B = VBMGxt[1]
    M = VBMGxt[2]
    G = VBMGxt[3]
    x_totalPoint = VBMGxt.shape[1]
    # there are n dSdt
    dM_dt_array = np.zeros(x_totalPoint)
    # each dSdt with the same equation form
    dM_dt_array[:] = +inRateM*B[:] - consumeRateM*M[:]*V[:] - outRateM*M[:]
    return(dM_dt_array)

def dGdt_array(VBMGxt = [], *args):
    # naming
    V = VBMGxt[0]
    B = VBMGxt[1]
    M = VBMGxt[2]
    G = VBMGxt[3]
    x_totalPoint = VBMGxt.shape[1]
    # there are n dSdt
    dG_dt_array = np.zeros(x_totalPoint)
    # each dSdt with the same equation form
    Gcopy = np.copy(G)
    centerX = Gcopy[:]
    leftX = np.roll(Gcopy[:], 1)
    rightX = np.roll(Gcopy[:], -1)
    leftX[0] = centerX[0]
    rightX[-1] = centerX[-1]
    dG_dt_array[:] = +(inRateG + alva.event_OAS)*B[:] - consumeRateG*G[:]*V[:] - outRateG*G[:]                      + mutatRateA*(leftX[:] - 2*centerX[:] + rightX[:])/(dx**2)
    return(dG_dt_array)

# In[7]:

# setting parameter
timeUnit = 'day'
if timeUnit == 'hour':
    hour = float(1)
    day = float(24)
elif timeUnit == 'day':
    day = float(1)
    hour = float(1)/24
elif timeUnit == 'year':
    year = float(1)
    day = float(1)/365
    hour = float(1)/24/365

maxV = float(50) # max virus/micro-liter
inRateV = 0.2/hour # in-rate of virus
killRateVm = 0.0003/hour # kill-rate of virus by antibody-IgM
killRateVg = killRateVm # kill-rate of virus by antibody-IgG

inRateB = 0.06/hour # in-rate of B-cell
outRateB = inRateB/8 # out-rate of B-cell
actRateBm = killRateVm # activation rate of naive B-cell

inRateM = 0.16/hour # in-rate of antibody-IgM from naive B-cell
outRateM = inRateM/1  # out-rate of antibody-IgM from naive B-cell
consumeRateM = killRateVm # consume-rate of antibody-IgM by cleaning virus

inRateG = inRateM/10 # in-rate of antibody-IgG from memory B-cell
outRateG = outRateM/250 # out-rate of antibody-IgG from memory B-cell
consumeRateG = killRateVg  # consume-rate of antibody-IgG by cleaning virus

mutatRateB = 0.00003/hour # B-cell mutation rate
mutatRateA = 0.0001/hour # antibody mutation rate

mutatRateB = 0.0000/hour # B-cell mutation rate
mutatRateA = 0.000/hour # antibody mutation rate

# time boundary and griding condition
minT = float(0)
maxT = float(6*28*day)
totalPoint_T = int(1*10**3 + 1)
gT = np.linspace(minT, maxT, totalPoint_T)
spacingT = np.linspace(minT, maxT, num = totalPoint_T, retstep = True)
gT = spacingT[0]
dt = spacingT[1]

# space boundary and griding condition
minX = float(0)
maxX = float(3)
totalPoint_X = int(maxX - minX + 1)
gX = np.linspace(minX, maxX, totalPoint_X)
gridingX = np.linspace(minX, maxX, num = totalPoint_X, retstep = True)
gX = gridingX[0]
dx = gridingX[1]
gV_array = np.zeros([totalPoint_X, totalPoint_T])
gB_array = np.zeros([totalPoint_X, totalPoint_T])
gM_array = np.zeros([totalPoint_X, totalPoint_T])
gG_array = np.zeros([totalPoint_X, totalPoint_T])
# initial output condition
#gV_array[1, 0] = float(2)
#[pre-parameter, post-parameter, recovered-day, OAS+, OSA-]
actRateBg_1st = 0.0002/hour # activation rate of memory B-cell at 1st time (pre-)
actRateBg_2nd = actRateBg_1st*10 # activation rate of memory B-cell at 2nd time (post-)
origin_virus = int(1)
current_virus = int(2)
event_parameter = np.array([[actRateBg_1st,
                            actRateBg_2nd,
                            14*day,
                            +5/hour,
                            -actRateBm - actRateBg_1st + (actRateBm + actRateBg_1st)/3,
                            origin_virus,
                            current_virus]])
# [viral population, starting time, first]
# [viral population, starting time] ---first
infection_period = 1*28*day
viral_population = np.zeros(int(maxX + 1))
viral_population[origin_virus:current_virus + 1] = 3
infection_starting_time = np.arange(int(maxX + 1))*infection_period
event_1st = np.zeros([int(maxX + 1), 2])
event_1st[:, 0] = viral_population
event_1st[:, 1] = infection_starting_time
print ('event_1st = {:}'.format(event_1st))

# [viral population, starting time] ---2nd]
viral_population = np.zeros(int(maxX + 1))
viral_population[origin_virus:current_virus + 1] = 0
infection_starting_time = np.arange(int(maxX + 1))*0
event_2nd = np.zeros([int(maxX + 1), 2])
event_2nd[:, 0] = viral_population
event_2nd[:, 1] = infection_starting_time
print ('event_2nd = {:}'.format(event_2nd))

event_table = np.array([event_parameter, event_1st, event_2nd])

# Runge Kutta numerical solution
pde_array = np.array([dVdt_array, dBdt_array, dMdt_array, dGdt_array])
initial_Out = np.array([gV_array, gB_array, gM_array, gG_array])
gOut_array = alva.AlvaRungeKutta4XT(pde_array, initial_Out, minX, maxX, totalPoint_X, minT, maxT, totalPoint_T, event_table)

# plotting
gV = gOut_array[0]
gB = gOut_array[1]
gM = gOut_array[2]
gG = gOut_array[3]

numberingFig = numberingFig + 1
for i in range(totalPoint_X):
    figure_name = '-response-%i'%(i)
    figure_suffix = '.png'
    save_figure = os.path.join(dir_path, file_name + figure_name + file_suffix)
    plt.figure(numberingFig, figsize = AlvaFigSize)
    plt.plot(gT, gV[i], color = 'red', label = r'$ V_{%i}(t) $'%(i), linewidth = 3.0, alpha = 0.5)
    plt.plot(gT, gM[i], color = 'blue', label = r'$ IgM_{%i}(t) $'%(i), linewidth = 3.0, alpha = 0.5)
    plt.plot(gT, gG[i], color = 'green', label = r'$ IgG_{%i}(t) $'%(i), linewidth = 3.0, alpha = 0.5)
    plt.plot(gT, gM[i] + gG[i], color = 'gray', linewidth = 5.0, alpha = 0.5, linestyle = 'dashed'
             , label = r'$ IgM_{%i}(t) + IgG_{%i}(t) $'%(i, i))
    plt.grid(True, which = 'both')
    plt.title(r'$ Antibody \ from \ Virus-{%i} $'%(i), fontsize = AlvaFontSize)
    plt.xlabel(r'$time \ (%s)$'%(timeUnit), fontsize = AlvaFontSize)
    plt.ylabel(r'$ Neutralization \ \ titer $', fontsize = AlvaFontSize)
    plt.xlim([minT, maxT])
    plt.xticks(fontsize = AlvaFontSize*0.6)
    plt.yticks(fontsize = AlvaFontSize*0.6)
    plt.ylim([2**0, 2**14])
    plt.yscale('log', basey = 2)
    plt.legend(loc = (1,0), fontsize = AlvaFontSize)
    plt.savefig(save_figure, dpi = 100)
    plt.show()

# In[5]:

# Experimental lab data from OAS paper
gT_lab = np.array([28, 28 + 7, 28 + 14, 28 + 28]) + 28
gPR8_lab = np.array([2**(9 + 1.0/10), 2**(13 - 1.0/5), 2**(13 + 1.0/3), 2**(13 - 1.0/4)])
standard_PR8 = gPR8_lab**(3.0/4)

gFM1_lab = np.array([0, 2**(6 - 1.0/5), 2**(7 - 1.0/4), 2**(8 + 1.0/4)])
standard_FM1 = gFM1_lab**(3.0/4)
bar_width = 2.0

# Sequential immunization graph

numberingFig = numberingFig + 1
plt.figure(numberingFig, figsize = (12, 6))
plt.subplot(111)
plt.plot(gT, (gM[origin_virus] + gG[origin_virus]), linewidth = 5.0, alpha = 0.5, color = 'gray'
         , label = r'$ Origin-virus $')
plt.plot(gT, (gM[origin_virus + 1] + gG[origin_virus + 1]), linewidth = 5.0, alpha = 0.5, color = 'red'
         , label = r'$ Subsequence-virus $')
plt.bar(gT_lab - bar_width/2, gPR8_lab, bar_width, alpha = 0.6, color = 'gray', yerr = standard_PR8
        , error_kw = dict(elinewidth = 1, ecolor = 'black'), label = r'$ PR8-virus $')
plt.bar(gT_lab + bar_width/2, gFM1_lab, bar_width, alpha = 0.6, color = 'red', yerr = standard_FM1
        , error_kw = dict(elinewidth = 1, ecolor = 'black'), label = r'$ FM1-virus $')
plt.grid(True, which = 'both')
plt.title(r'$ Original \ Antigenic \ Sin \ (sequential-infection)$', fontsize = AlvaFontSize)
plt.xlabel(r'$time \ (%s)$'%(timeUnit), fontsize = AlvaFontSize)
plt.ylabel(r'$ Neutralization \ \ titer $', fontsize = AlvaFontSize)
plt.xticks(fontsize = AlvaFontSize*0.6)
plt.yticks(fontsize = AlvaFontSize*0.6)
plt.xlim([minT, 6*30*day])
plt.ylim([2**5, 2**14])
plt.yscale('log', basey = 2)
# gca()---GetCurrentAxis and Format the ticklabel to be 2**x
plt.gca().yaxis.set_major_formatter(FuncFormatter(lambda x, pos: int(2**(np.log(x)/np.log(2)))))
#plt.gca().xaxis.set_major_locator(plt.MultipleLocator(7))
plt.legend(loc = (1, 0), fontsize = AlvaFontSize)
plt.show()

# In[6]:

# Experimental lab data from OAS paper
gT_lab = np.array([28, 28 + 7, 28 + 14, 28 + 28]) + 28
gPR8_lab = np.array([2**(9 + 1.0/10), 2**(13 - 1.0/5), 2**(13 + 1.0/3), 2**(13 - 1.0/4)])
standard_PR8 = gPR8_lab**(3.0/4)

gFM1_lab = np.array([0, 2**(6 - 1.0/5), 2**(7 - 1.0/4), 2**(8 + 1.0/4)])
standard_FM1 = gFM1_lab**(3.0/4)
bar_width = 1.0

# Sequential immunization graph
figure_name = '-Original-Antigenic-Sin-infection'
figure_suffix = '.png'
save_figure = os.path.join(dir_path, file_name + figure_name + file_suffix)
numberingFig = numberingFig + 1
plt.figure(numberingFig, figsize = (12, 6))
plt.subplot(111)
plt.plot(gT, (gM[origin_virus] + gG[origin_virus]), linewidth = 5.0, alpha = 0.5, color = 'gray'
         , label = r'$ Origin-virus $')
plt.plot(gT, (gM[origin_virus + 1] + gG[origin_virus + 1]), linewidth = 5.0, alpha = 0.5, color = 'red'
         , label = r'$ Subsequence-virus $')
plt.bar(gT_lab - bar_width/2, gPR8_lab, bar_width, alpha = 0.6, color = 'gray', yerr = standard_PR8
        , error_kw = dict(elinewidth = 1, ecolor = 'black'), label = r'$ PR8-virus $')
plt.bar(gT_lab + bar_width/2, gFM1_lab, bar_width, alpha = 0.6, color = 'red', yerr = standard_FM1
        , error_kw = dict(elinewidth = 1, ecolor = 'black'), label = r'$ FM1-virus $')
plt.grid(True, which = 'both')
plt.title(r'$ Original \ Antigenic \ Sin \ (sequential-infection)$', fontsize = AlvaFontSize)
plt.xlabel(r'$time \ (%s)$'%(timeUnit), fontsize = AlvaFontSize)
plt.ylabel(r'$ Neutralization \ \ titer $', fontsize = AlvaFontSize)
plt.xticks(fontsize = AlvaFontSize*0.6)
plt.yticks(fontsize = AlvaFontSize*0.6)
plt.xlim([minT, 3*30*day])
plt.ylim([2**5, 2**14])
plt.yscale('log', basey = 2)
# gca()---GetCurrentAxis and Format the ticklabel to be 2**x
plt.gca().yaxis.set_major_formatter(FuncFormatter(lambda x, pos: int(2**(np.log(x)/np.log(2)))))
plt.gca().xaxis.set_major_locator(plt.MultipleLocator(7))
plt.legend(loc = (1, 0), fontsize = AlvaFontSize)
plt.savefig(save_figure, dpi = 100, bbox_inches='tight')
plt.show()

# In[ ]:


#!/usr/bin/python2.4
import numpy
import scipy.stats
import matplotlib.pyplot as plt

def main():
    # true parameters
    w  = 2
    w0 = 3
    sigma = 2

    # make data
    numpy.random.seed(1)
    Ntrain = 20
    xtrain = numpy.linspace(0,10,Ntrain)
    ytrain = w*xtrain + w0 + numpy.random.random(Ntrain)*sigma
    Ntest = 100
    xtest = numpy.linspace(0,10,Ntest)
    ytest = w*xtest + w0 + numpy.random.random(Ntest)*sigma

    #  from http://www2.warwick.ac.uk/fac/sci/moac/students/peter_cock/python/lin_reg/
    # fit
    west, w0est, r_value, p_value, std_err = scipy.stats.linregress(xtrain, ytrain)

    # display
    print "Param \t True \t Est"
    print "w0 \t %5.3f \t %5.3f" % (w0, w0est)
    print "w \t %5.3f \t %5.3f" % (w, west)

    # plot
    plt.close()
    plt.plot(xtrain, ytrain, 'ro')
    plt.hold(True)
    #plt.plot(xtest, ytest, 'ka-')
    ytestPred = west*xtest + w0est
    #ndx = range(0, Ntest, 10)
    #h = plt.plot(xtest[ndx], ytestPred[ndx], 'b*')
    h = plt.plot(xtest, ytestPred, 'b-')
    plt.setp(h, 'markersize', 12)

if __name__ == '__main__':
	main()


#!/usr/bin/env python
#
# Copyright 2013 Free Software Foundation, Inc.
#
# This file is part of GNU Radio
#
# GNU Radio is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3, or (at your option)
# any later version.
#
# GNU Radio is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with GNU Radio; see the file COPYING.  If not, write to
# the Free Software Foundation, Inc., 51 Franklin Street,
# Boston, MA 02110-1301, USA.
#

from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals
from gnuradio import gr, filter
from gnuradio import analog
from gnuradio import blocks
from gnuradio import eng_notation
from gnuradio.eng_arg import eng_float, intx
from argparse import ArgumentParser
import sys
import numpy

try:
    from matplotlib import pyplot
except ImportError:
    print("Error: could not from matplotlib import pyplot (http://matplotlib.sourceforge.net/)")
    sys.exit(1)

class example_fft_filter_ccc(gr.top_block):
    def __init__(self, N, fs, bw0, bw1, tw, atten, D):
        gr.top_block.__init__(self)

        self._nsamps = N
        self._fs = fs
        self._bw0 = bw0
        self._bw1 = bw1
        self._tw = tw
        self._at = atten
        self._decim = D
        taps = filter.firdes.complex_band_pass_2(1, self._fs,
                                                 self._bw0, self._bw1,
                                                 self._tw, self._at)
        print("Num. Taps: ", len(taps))

        self.src  = analog.noise_source_c(analog.GR_GAUSSIAN, 1)
        self.head = blocks.head(gr.sizeof_gr_complex, self._nsamps)

        self.filt0 = filter.fft_filter_ccc(self._decim, taps)

        self.vsnk_src = blocks.vector_sink_c()
        self.vsnk_out = blocks.vector_sink_c()

        self.connect(self.src, self.head, self.vsnk_src)
        self.connect(self.head, self.filt0, self.vsnk_out)

def main():
    parser = ArgumentParser(conflict_handler="resolve")
    parser.add_argument("-N", "--nsamples", type=int, default=10000,
                      help="Number of samples to process [default=%(default)r]")
    parser.add_argument("-s", "--samplerate", type=eng_float, default=8000,
                      help="System sample rate [default=%(default)r]")
    parser.add_argument("-S", "--start-pass", type=eng_float, default=1000,
                      help="Start of Passband [default=%(default)r]")
    parser.add_argument("-E", "--end-pass", type=eng_float, default=2000,
                      help="End of Passband [default=%(default)r]")
    parser.add_argument("-T", "--transition", type=eng_float, default=100,
                      help="Transition band [default=%(default)r]")
    parser.add_argument("-A", "--attenuation", type=eng_float, default=80,
                      help="Stopband attenuation [default=%(default)r]")
    parser.add_argument("-D", "--decimation", type=int, default=1,
                      help="Decmation factor [default=%(default)r]")
    args = parser.parse_args()

    put = example_fft_filter_ccc(args.nsamples,
                                 args.samplerate,
                                 args.start_pass,
                                 args.end_pass,
                                 args.transition,
                                 args.attenuation,
                                 args.decimation)
    put.run()

    data_src = numpy.array(put.vsnk_src.data())
    data_snk = numpy.array(put.vsnk_out.data())

    # Plot the signals PSDs
    nfft = 1024
    f1 = pyplot.figure(1, figsize=(12,10))
    s1 = f1.add_subplot(1,1,1)
    s1.psd(data_src, NFFT=nfft, noverlap=nfft / 4,
           Fs=args.samplerate)
    s1.psd(data_snk, NFFT=nfft, noverlap=nfft / 4,
           Fs=args.samplerate)

    f2 = pyplot.figure(2, figsize=(12,10))
    s2 = f2.add_subplot(1,1,1)
    s2.plot(data_src)
    s2.plot(data_snk.real, 'g')

    pyplot.show()

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass


"""
The :mod:`sklearn.lda` module implements Linear Discriminant Analysis (LDA).
"""
# Authors: Matthieu Perrot
#          Mathieu Blondel

import warnings

import numpy as np
from scipy import linalg

from .base import BaseEstimator, ClassifierMixin, TransformerMixin
from .utils.extmath import logsumexp
from .utils.fixes import unique
from .utils import check_arrays

__all__ = ['LDA']

class LDA(BaseEstimator, ClassifierMixin, TransformerMixin):
    """
    Linear Discriminant Analysis (LDA)

    A classifier with a linear decision boundary, generated
    by fitting class conditional densities to the data
    and using Bayes' rule.

    The model fits a Gaussian density to each class, assuming that
    all classes share the same covariance matrix.

    The fitted model can also be used to reduce the dimensionality
    of the input, by projecting it to the most discriminative
    directions.

    Parameters
    ----------

    n_components: int
        Number of components (< n_classes - 1) for dimensionality reduction

    priors : array, optional, shape = [n_classes]
        Priors on classes

    Attributes
    ----------
    `means_` : array-like, shape = [n_classes, n_features]
        Class means
    `xbar_` : float, shape = [n_features]
        Over all mean
    `priors_` : array-like, shape = [n_classes]
        Class priors (sum to 1)
    `covariance_` : array-like, shape = [n_features, n_features]
        Covariance matrix (shared by all classes)

    Examples
    --------
    >>> import numpy as np
    >>> from sklearn.lda import LDA
    >>> X = np.array([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]])
    >>> y = np.array([1, 1, 1, 2, 2, 2])
    >>> clf = LDA()
    >>> clf.fit(X, y)
    LDA(n_components=None, priors=None)
    >>> print(clf.predict([[-0.8, -1]]))
    [1]

    See also
    --------
    sklearn.qda.QDA: Quadratic discriminant analysis

    """

    def __init__(self, n_components=None, priors=None):
        self.n_components = n_components
        self.priors = np.asarray(priors) if priors is not None else None

        if self.priors is not None:
            if (self.priors < 0).any():
                raise ValueError('priors must be non-negative')
            if self.priors.sum() != 1:
                print 'warning: the priors do not sum to 1. Renormalizing'
                self.priors = self.priors / self.priors.sum()

    def fit(self, X, y, store_covariance=False, tol=1.0e-4):
        """
        Fit the LDA model according to the given training data and parameters.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
            Training vector, where n_samples in the number of samples and
            n_features is the number of features.
        y : array, shape = [n_samples]
            Target values (integers)
        store_covariance : boolean
            If True the covariance matrix (shared by all classes) is computed
            and stored in `self.covariance_` attribute.
        """
        X, y = check_arrays(X, y, sparse_format='dense')
        self.classes_, y = unique(y, return_inverse=True)
        n_samples, n_features = X.shape
        n_classes = len(self.classes_)
        if n_classes < 2:
            raise ValueError('y has less than 2 classes')
        if self.priors is None:
            self.priors_ = np.bincount(y) / float(n_samples)
        else:
            self.priors_ = self.priors

        # Group means n_classes*n_features matrix
        means = []
        Xc = []
        cov = None
        if store_covariance:
            cov = np.zeros((n_features, n_features))
        for ind in xrange(n_classes):
            Xg = X[y == ind, :]
            meang = Xg.mean(0)
            means.append(meang)
            # centered group data
            Xgc = Xg - meang
            Xc.append(Xgc)
            if store_covariance:
                cov += np.dot(Xgc.T, Xgc)
        if store_covariance:
            cov /= (n_samples - n_classes)
            self.covariance_ = cov

        self.means_ = np.asarray(means)
        Xc = np.concatenate(Xc, 0)

        # ----------------------------
        # 1) within (univariate) scaling by with classes std-dev
        std = Xc.std(axis=0)
        # avoid division by zero in normalization
        std[std == 0] = 1.
        fac = float(1) / (n_samples - n_classes)
        # ----------------------------
        # 2) Within variance scaling
        X = np.sqrt(fac) * (Xc / std)
        # SVD of centered (within)scaled data
        U, S, V = linalg.svd(X, full_matrices=0)

        rank = np.sum(S > tol)
        if rank < n_features:
            warnings.warn("Variables are collinear")
        # Scaling of within covariance is: V' 1/S
        scaling = (V[:rank] / std).T / S[:rank]

        ## ----------------------------
        ## 3) Between variance scaling
        # Overall mean
        xbar = np.dot(self.priors_, self.means_)
        # Scale weighted centers
        X = np.dot(((np.sqrt((n_samples * self.priors_) * fac)) *
                    (means - xbar).T).T, scaling)
        # Centers are living in a space with n_classes-1 dim (maximum)
        # Use svd to find projection in the space spanned by the
        # (n_classes) centers
        _, S, V = linalg.svd(X, full_matrices=0)

        rank = np.sum(S > tol * S[0])
        # compose the scalings
        self.scaling = np.dot(scaling, V.T[:, :rank])
        self.xbar_ = xbar
        # weight vectors / centroids
        self.coef_ = np.dot(self.means_ - self.xbar_, self.scaling)
        self.intercept_ = -0.5 * np.sum(self.coef_ ** 2, axis=1) + \
                           np.log(self.priors_)
        return self

    @property
    def classes(self):
        warnings.warn("LDA.classes is deprecated and will be removed in 0.14. "
                      "Use LDA.classes_ instead.", DeprecationWarning,
                      stacklevel=2)
        return self.classes_

    def _decision_function(self, X):
        X = np.asarray(X)
        # center and scale data
        X = np.dot(X - self.xbar_, self.scaling)
        return np.dot(X, self.coef_.T) + self.intercept_

    def decision_function(self, X):
        """
        This function return the decision function values related to each
        class on an array of test vectors X.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]

        Returns
        -------
        C : array, shape = [n_samples, n_classes] or [n_samples,]
            Decision function values related to each class, per sample.
            In the two-class case, the shape is [n_samples,], giving the
            log likelihood ratio of the positive class.
        """
        dec_func = self._decision_function(X)
        if len(self.classes_) == 2:
            return dec_func[:, 1] - dec_func[:, 0]
        return dec_func

    def transform(self, X):
        """
        Project the data so as to maximize class separation (large separation
        between projected class means and small variance within each class).

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]

        Returns
        -------
        X_new : array, shape = [n_samples, n_components]
        """
        X = np.asarray(X)
        # center and scale data
        X = np.dot(X - self.xbar_, self.scaling)
        n_comp = X.shape[1] if self.n_components is None else self.n_components
        return np.dot(X, self.coef_[:n_comp].T)

    def predict(self, X):
        """
        This function does classification on an array of test vectors X.

        The predicted class C for each sample in X is returned.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]

        Returns
        -------
        C : array, shape = [n_samples]
        """
        d = self._decision_function(X)
        y_pred = self.classes_.take(d.argmax(1))
        return y_pred

    def predict_proba(self, X):
        """
        This function return posterior probabilities of classification
        according to each class on an array of test vectors X.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]

        Returns
        -------
        C : array, shape = [n_samples, n_classes]
        """
        values = self._decision_function(X)
        # compute the likelihood of the underlying gaussian models
        # up to a multiplicative constant.
        likelihood = np.exp(values - values.max(axis=1)[:, np.newaxis])
        # compute posterior probabilities
        return likelihood / likelihood.sum(axis=1)[:, np.newaxis]

    def predict_log_proba(self, X):
        """
        This function return posterior log-probabilities of classification
        according to each class on an array of test vectors X.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]

        Returns
        -------
        C : array, shape = [n_samples, n_classes]
        """
        values = self._decision_function(X)
        loglikelihood = (values - values.max(axis=1)[:, np.newaxis])
        normalization = logsumexp(loglikelihood, axis=1)
        return loglikelihood - normalization[:, np.newaxis]


#!/usr/bin/env python

# Author: Shao Zhang, Phil Saltzman, and Eddie Canaan
# Last Updated: 2015-03-13
#
# This tutorial will demonstrate some uses for intervals in Panda
# to move objects in your panda world.
# Intervals are tools that change a value of something, like position,
# rotation or anything else, linearly, over a set period of time. They can be
# also be combined to work in sequence or in Parallel
#
# In this lesson, we will simulate a carousel in motion using intervals.
# The carousel will spin using an hprInterval while 4 pandas will represent
# the horses on a traditional carousel. The 4 pandas will rotate with the
# carousel and also move up and down on their poles using a LerpFunc interval.
# Finally there will also be lights on the outer edge of the carousel that
# will turn on and off by switching their texture with intervals in Sequence
# and Parallel

from direct.showbase.ShowBase import ShowBase
from panda3d.core import AmbientLight, DirectionalLight, LightAttrib
from panda3d.core import NodePath
from panda3d.core import LVector3
from direct.interval.IntervalGlobal import *  # Needed to use Intervals
from direct.gui.DirectGui import *

# Importing math constants and functions
from math import pi, sin

class CarouselDemo(ShowBase):

    def __init__(self):
        # Initialize the ShowBase class from which we inherit, which will
        # create a window and set up everything we need for rendering into it.
        ShowBase.__init__(self)

        # This creates the on screen title that is in every tutorial
        self.title = OnscreenText(text="Panda3D: Tutorial - Carousel",
                                  parent=base.a2dBottomCenter,
                                  fg=(1, 1, 1, 1), shadow=(0, 0, 0, .5),
                                  pos=(0, .1), scale=.1)

        base.disableMouse()  # Allow manual positioning of the camera
        camera.setPosHpr(0, -8, 2.5, 0, -9, 0)  # Set the cameras' position
                                                # and orientation

        self.loadModels()  # Load and position our models
        self.setupLights()  # Add some basic lighting
        self.startCarousel()  # Create the needed intervals and put the
                              # carousel into motion

    def loadModels(self):
        # Load the carousel base
        self.carousel = loader.loadModel("models/carousel_base")
        self.carousel.reparentTo(render)  # Attach it to render

        # Load the modeled lights that are on the outer rim of the carousel
        # (not Panda lights)
        # There are 2 groups of lights. At any given time, one group will have
        # the "on" texture and the other will have the "off" texture.
        self.lights1 = loader.loadModel("models/carousel_lights")
        self.lights1.reparentTo(self.carousel)

        # Load the 2nd set of lights
        self.lights2 = loader.loadModel("models/carousel_lights")
        # We need to rotate the 2nd so it doesn't overlap with the 1st set.
        self.lights2.setH(36)
        self.lights2.reparentTo(self.carousel)

        # Load the textures for the lights. One texture is for the "on" state,
        # the other is for the "off" state.
        self.lightOffTex = loader.loadTexture("models/carousel_lights_off.jpg")
        self.lightOnTex = loader.loadTexture("models/carousel_lights_on.jpg")

        # Create an list (self.pandas) with filled with 4 dummy nodes attached
        # to the carousel.
        # This uses a python concept called "Array Comprehensions."  Check the
        # Python manual for more information on how they work
        self.pandas = [self.carousel.attachNewNode("panda" + str(i))
                       for i in range(4)]
        self.models = [loader.loadModel("models/carousel_panda")
                       for i in range(4)]
        self.moves = [0] * 4

        for i in range(4):
            # set the position and orientation of the ith panda node we just created
            # The Z value of the position will be the base height of the pandas.
            # The headings are multiplied by i to put each panda in its own position
            # around the carousel
            self.pandas[i].setPosHpr(0, 0, 1.3, i * 90, 0, 0)

            # Load the actual panda model, and parent it to its dummy node
            self.models[i].reparentTo(self.pandas[i])
            # Set the distance from the center. This distance is based on the way the
            # carousel was modeled in Maya
            self.models[i].setY(.85)

        # Load the environment (Sky sphere and ground plane)
        self.env = loader.loadModel("models/env")
        self.env.reparentTo(render)
        self.env.setScale(7)

    # Panda Lighting
    def setupLights(self):
        # Create some lights and add them to the scene. By setting the lights on
        # render they affect the entire scene
        # Check out the lighting tutorial for more information on lights
        ambientLight = AmbientLight("ambientLight")
        ambientLight.setColor((.4, .4, .35, 1))
        directionalLight = DirectionalLight("directionalLight")
        directionalLight.setDirection(LVector3(0, 8, -2.5))
        directionalLight.setColor((0.9, 0.8, 0.9, 1))
        render.setLight(render.attachNewNode(directionalLight))
        render.setLight(render.attachNewNode(ambientLight))

        # Explicitly set the environment to not be lit
        self.env.setLightOff()

    def startCarousel(self):
        # Here's where we actually create the intervals to move the carousel
        # The first type of interval we use is one created directly from a NodePath
        # This interval tells the NodePath to vary its orientation (hpr) from its
        # current value (0,0,0) to (360,0,0) over 20 seconds. Intervals created from
        # NodePaths also exist for position, scale, color, and shear

        self.carouselSpin = self.carousel.hprInterval(20, LVector3(360, 0, 0))
        # Once an interval is created, we need to tell it to actually move.
        # start() will cause an interval to play once. loop() will tell an interval
        # to repeat once it finished. To keep the carousel turning, we use
        # loop()
        self.carouselSpin.loop()

        # The next type of interval we use is called a LerpFunc interval. It is
        # called that becuase it linearly interpolates (aka Lerp) values passed to
        # a function over a given amount of time.

        # In this specific case, horses on a carousel don't move contantly up,
        # suddenly stop, and then contantly move down again. Instead, they start
        # slowly, get fast in the middle, and slow down at the top. This motion is
        # close to a sine wave. This LerpFunc calls the function oscillatePanda
        # (which we will create below), which changes the height of the panda based
        # on the sin of the value passed in. In this way we achieve non-linear
        # motion by linearly changing the input to a function
        for i in range(4):
            self.moves[i] = LerpFunc(
                self.oscillatePanda,  # function to call
                duration=3,  # 3 second duration
                fromData=0,  # starting value (in radians)
                toData=2 * pi,  # ending value (2pi radians = 360 degrees)
                # Additional information to pass to
                # self.oscialtePanda
                extraArgs=[self.models[i], pi * (i % 2)]
            )
            # again, we want these to play continuously so we start them with
            # loop()
            self.moves[i].loop()

        # Finally, we combine Sequence, Parallel, Func, and Wait intervals,
        # to schedule texture swapping on the lights to simulate the lights turning
        # on and off.
        # Sequence intervals play other intervals in a sequence. In other words,
        # it waits for the current interval to finish before playing the next
        # one.
        # Parallel intervals play a group of intervals at the same time
        # Wait intervals simply do nothing for a given amount of time
        # Func intervals simply make a single function call. This is helpful because
        # it allows us to schedule functions to be called in a larger sequence. They
        # take virtually no time so they don't cause a Sequence to wait.

        self.lightBlink = Sequence(
            # For the first step in our sequence we will set the on texture on one
            # light and set the off texture on the other light at the same time
            Parallel(
                Func(self.lights1.setTexture, self.lightOnTex, 1),
                Func(self.lights2.setTexture, self.lightOffTex, 1)),
            Wait(1),  # Then we will wait 1 second
            # Then we will switch the textures at the same time
            Parallel(
                Func(self.lights1.setTexture, self.lightOffTex, 1),
                Func(self.lights2.setTexture, self.lightOnTex, 1)),
            Wait(1)  # Then we will wait another second
        )

        self.lightBlink.loop()  # Loop this sequence continuously

    def oscillatePanda(self, rad, panda, offset):
        # This is the oscillation function mentioned earlier. It takes in a
        # degree value, a NodePath to set the height on, and an offset. The
        # offset is there so that the different pandas can move opposite to
        # each other.  The .2 is the amplitude, so the height of the panda will
        # vary from -.2 to .2
        panda.setZ(sin(rad + offset) * .2)

demo = CarouselDemo()
demo.run()