# Copyright (C) 2016 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import numpy as np


# call payoff for basket
class BasketCall :
    
    m_strike = 0.

    # Contructor
    # p_strike   strike
    def __init__(self, p_strike) :
        
        self.m_strike = p_strike
    
    # for one simulation, give the basket value
    # p_assets   For the current simulation values of all assets
    def apply(self, p_assets) :
        
        return np.maximum(np.mean(p_assets) - self.m_strike, np.zeros(len(p_assets.transpose())))

    # for all simulations, give the basket values
    # p_assets    values of all assets (one column is a simulation)
    def applyVec(self, p_assets) :
        
        ret = np.zeros(p_assets.shape[1])
        
        return np.maximum(np.mean(p_assets, axis=0) - self.m_strike, np.zeros(len(p_assets.transpose())))
        
# Put payoff for basket
class BasketPut :
    
    m_strike = 0.

    # Constructor
    # p_strike   strike
    def __init__(self, p_strike) :
        
        self.m_strike = p_strike

    # for all simulations, give the basket values
    # p_assets    values of all assets (one column is a simulation)
    def operator(self, p_assets) :
        
        ret = np.zeros(p_assets.shape[1])
        
        return np.maximum(-np.mean(p_assets, axis=0) + self.m_strike, np.zeros(len(p_assets.transpose())))