from scipy.stats import poisson
import math

class ModelInfo:

    def __init__(self,ARRIVAL_RATE,STARTING_INVENTORY,CUSTOMER_1_PROP,CUSTOMER_2_PROP,
                 HIGH_PRICE,LOW_PRICE) -> None:
        self.ARRIVAL_RATE =ARRIVAL_RATE
        self.STARTING_INVENTORY = STARTING_INVENTORY
        self.CUSTOMER_1_PROP = CUSTOMER_1_PROP
        self.CUSTOMER_2_PROP = CUSTOMER_2_PROP
        self.HIGH_PRICE = HIGH_PRICE
        self.LOW_PRICE = LOW_PRICE

    def set_ARRIVAL_RATE(self,ARRIVAL_RATE):
        return ModelInfo(ARRIVAL_RATE=ARRIVAL_RATE,
                         STARTING_INVENTORY=self.STARTING_INVENTORY,
                         CUSTOMER_1_PROP=self.CUSTOMER_1_PROP,
                         CUSTOMER_2_PROP=self.CUSTOMER_2_PROP,
                         HIGH_PRICE=self.HIGH_PRICE,
                         LOW_PRICE=self.LOW_PRICE)

    def set_STARTING_INVENTORY(self,STARTING_INVENTORY):
        return ModelInfo(ARRIVAL_RATE=self.ARRIVAL_RATE,
                         STARTING_INVENTORY=STARTING_INVENTORY,
                         CUSTOMER_1_PROP=self.CUSTOMER_1_PROP,
                         CUSTOMER_2_PROP=self.CUSTOMER_2_PROP,
                         HIGH_PRICE=self.HIGH_PRICE,
                         LOW_PRICE=self.LOW_PRICE)
    
    def set_CUSTOMER_1_PROP(self,CUSTOMER_1_PROP):
        return ModelInfo(ARRIVAL_RATE=self.ARRIVAL_RATE,
                         STARTING_INVENTORY=self.STARTING_INVENTORY,
                         CUSTOMER_1_PROP=CUSTOMER_1_PROP,
                         CUSTOMER_2_PROP=self.CUSTOMER_2_PROP,
                         HIGH_PRICE=self.HIGH_PRICE,
                         LOW_PRICE=self.LOW_PRICE)
    
    def set_CUSTOMER_2_PROP(self,CUSTOMER_2_PROP):
        return ModelInfo(ARRIVAL_RATE=self.ARRIVAL_RATE,
                         STARTING_INVENTORY=self.STARTING_INVENTORY,
                         CUSTOMER_1_PROP=self.CUSTOMER_1_PROP,
                         CUSTOMER_2_PROP=CUSTOMER_2_PROP,
                         HIGH_PRICE=self.HIGH_PRICE,
                         LOW_PRICE=self.LOW_PRICE)
    
    def set_HIGH_PRICE(self,HIGH_PRICE):
        return ModelInfo(ARRIVAL_RATE=self.ARRIVAL_RATE,
                         STARTING_INVENTORY=self.STARTING_INVENTORY,
                         CUSTOMER_1_PROP=self.CUSTOMER_1_PROP,
                         CUSTOMER_2_PROP=self.CUSTOMER_2_PROP,
                         HIGH_PRICE=HIGH_PRICE,
                         LOW_PRICE=self.LOW_PRICE)
    
    def set_LOW_PRICE(self,LOW_PRICE):
        return ModelInfo(ARRIVAL_RATE=self.ARRIVAL_RATE,
                         STARTING_INVENTORY=self.STARTING_INVENTORY,
                         CUSTOMER_1_PROP=self.CUSTOMER_1_PROP,
                         CUSTOMER_2_PROP=self.CUSTOMER_2_PROP,
                         HIGH_PRICE=self.HIGH_PRICE,
                         LOW_PRICE=LOW_PRICE)

def expected_revenue_in_period_1(model:ModelInfo):
    mu = model.ARRIVAL_RATE
    c1 = model.STARTING_INVENTORY
    customer_1_prop = model.CUSTOMER_1_PROP
    customer_2_prop = model.CUSTOMER_2_PROP
    p1 = model.HIGH_PRICE

    expected_sales = 0

    for i in range(c1):
        expected_sales += i * poisson.pmf(i,mu * customer_2_prop)
    expected_sales += c1 * (1 - poisson.cdf(c1-1,mu * customer_2_prop))

    return expected_sales * p1

def expected_revenue_in_period_2_given_price(c2,p2,residual_customer_1_period_1,residual_customer_2_period_1,model:ModelInfo):
    mu = model.ARRIVAL_RATE
    customer_1_prop = model.CUSTOMER_1_PROP
    customer_2_prop = model.CUSTOMER_2_PROP
    expected_sales = 0
    if residual_customer_2_period_1 > 0:
        # All are sold to customer 2 in period 1
        return 0
    
    # residual_customer_2_period_1 will always be 0, so not included here
    total_residual = residual_customer_1_period_1

    if p2 == model.LOW_PRICE:
        
        # if there is already enough customer before more arrivals willing to buy, so sell to them
        if total_residual >= c2:
            # everyone will buy at price = 1
            return c2 * p2
        else:
            # finding the expectation of min(# of customer + residual, c2)
            # expressed as min(# of customer, c2 - residual) + residual
            # makes sure c2 - residual is not negative
            expected_sales += total_residual
            for i in range(c2 - total_residual):
                # the arrival of cust 1 from period 1 is fixed so mu nvr add the proportion of cust 1 arrivals
                expected_sales += i * poisson.pmf(i,mu)
            expected_sales += (c2 - total_residual) * (1 - poisson.cdf(c2 - total_residual - 1,mu))
            
        return expected_sales * p2

        
    elif p2 == model.HIGH_PRICE:
        # only customer 2 buys
        # residual customer 2 should be 0, since p1=2 and all customer 2 who arrived in period 1 will be buy
        for i in range(c2):
            expected_sales += i * poisson.pmf(i,mu * customer_2_prop)
        expected_sales += c2 * (1 - poisson.cdf(c2 - 1, mu * customer_2_prop))
        return expected_sales * p2

    elif p2 > model.HIGH_PRICE or p2 <=0:
        return 0
    
def expected_revenue_in_period_2(c2,residual_customer_1_period_1,residual_customer_2_period_1,c2_threshold,model:ModelInfo):
    if c2 >= c2_threshold:
        p2 = model.LOW_PRICE
    else:
        p2 = model.HIGH_PRICE
    return expected_revenue_in_period_2_given_price(c2=c2,p2=p2,
                                                    residual_customer_1_period_1=residual_customer_1_period_1,
                                                    residual_customer_2_period_1=residual_customer_2_period_1,
                                                    model=model)

def calculate_probability(c2,residual_customer_1,residual_customer_2,model:ModelInfo):
    # calculate the probablity having c2 inventory, residual_customer_1, residual_customer_2 in period 2
    mu = model.ARRIVAL_RATE
    customer_1_prop = model.CUSTOMER_1_PROP
    customer_2_prop = model.CUSTOMER_2_PROP
    c1 = model.STARTING_INVENTORY

    customer_2_period_1 = c1 - c2 + residual_customer_2
    if c2 >= 0 and residual_customer_2 == 0:
        # case when customer 2 arrival in period 1 is less than or equal to inventory in period 1
        prob_1 = poisson.pmf(customer_2_period_1,mu * customer_2_prop)
        # all customer 1 arrival in period 1 becomes residual_customer_1 in period 2
        prob_2 = poisson.pmf(residual_customer_1,mu * customer_1_prop)
        #print(f'prob_1:{prob_1}|prob_2:{prob_2}')
        return prob_1 * prob_2
    elif c2 == 0 and residual_customer_2 > 0:
        # case when customer 2 arrival in period 1 is more than inventory in period 1
        prob_2 = poisson.pmf(residual_customer_1,mu * customer_1_prop)
        # prob that total customer 2 in period 1 is c1 + residual_customer_2, > c1
        prob_3 = poisson.pmf(c1+residual_customer_2,mu * customer_2_prop)
        #print(f'prob_2:{prob_2}|prob_3:{prob_3}')
        return prob_2 * prob_3
    elif c2 > 0 and residual_customer_2 > 0:
        # if there is inventory in period 2, there should not be any residual customer 2
        return 0



def calculate_expected_total_revenue(c2_threshold,model:ModelInfo):
    total_expected_revenue_in_period_2 = 0
    #inv = 0
    #for RC1 in range(100):
    #    for RC2 in range(100):
    #        total_expected_revenue_in_period_2 += calculate_probability(inv,RC1,RC2) * expected_revenue_in_period_2(inv,RC1,RC2,c2_threshold=c2_threshold)
    #print(total_expected_revenue_in_period_2)

    RC2=0
    for inv in range(1,model.STARTING_INVENTORY + 1):
        # if there is no price change and remains high price, RC1 will have no effect, the expected revenue in period 2 
        # is solely dependent on the customer 2
        customer_2_period_1 = model.STARTING_INVENTORY - inv
        if inv < c2_threshold:
            # prob of no of customer 2 in period * expected_rev_in_period_2 given c2=inv,RC1 is any value (since all same),RC2=0
            total_expected_revenue_in_period_2 += poisson.pmf(customer_2_period_1,model.ARRIVAL_RATE * model.CUSTOMER_2_PROP) * expected_revenue_in_period_2(inv,0,0,c2_threshold,model=model)
        elif inv >= c2_threshold:
            #there is price change to LOW PRICE
            for RC1 in range(inv):
                # when residual is less than inv in period 2, there are still different values of prob and expected rev
                #inv=2 | RC1 =0,1 | RC2=0
                total_expected_revenue_in_period_2 += calculate_probability(inv,RC1,RC2,model=model) * expected_revenue_in_period_2(inv,RC1,RC2,c2_threshold=c2_threshold,model=model)
            #when residual is greater than or qual to inv in period 2, the expected rev is just inv * LOW_PRICE
            #inv =2 | RC1 >=2 | RC2=0 | expected rev in period 2 = inv * LOW_PIRCE = 2
            #prob of the no of cutomer 2 in period 1, same for all cases of RC1, thus we can factor it out
            prob_1 = poisson.pmf(customer_2_period_1,model.ARRIVAL_RATE * model.CUSTOMER_2_PROP)
            # prob_1 * ( prob_2 + prob_2 + .....)
            cum_prob_2 = 1-poisson.cdf(inv-1,model.ARRIVAL_RATE * model.CUSTOMER_1_PROP)
            total_expected_revenue_in_period_2 += prob_1 * cum_prob_2 * expected_revenue_in_period_2(inv,inv,0,c2_threshold,model=model)

    return total_expected_revenue_in_period_2 + expected_revenue_in_period_1(model=model)

def calculate_expected_total_revenue_given_low_low(model:ModelInfo):
    mu = 2*model.ARRIVAL_RATE
    c1 = model.STARTING_INVENTORY
    p = model.LOW_PRICE

    expected_sales = 0

    for i in range(c1):
        expected_sales += i * poisson.pmf(i,mu)
    expected_sales += c1 * (1 - poisson.cdf(c1-1,mu))

    return expected_sales * p



def calculate_expected_total_revenue_given_price(p2,model:ModelInfo):
    total_expected_revenue_in_period_2 = 0
    

    RC2=0
    for inv in range(1,model.STARTING_INVENTORY+1):
        customer_2_period_1 = model.STARTING_INVENTORY - inv
        for RC1 in range(inv):
            # RC1 =0,1
            total_expected_revenue_in_period_2 += calculate_probability(inv,RC1,RC2,model=model) * expected_revenue_in_period_2_given_price(inv,p2,RC1,RC2,model=model)
        #RC1 =2,3,...
        prob_1 = poisson.pmf(customer_2_period_1,model.ARRIVAL_RATE * model.CUSTOMER_2_PROP)
        # prob_1 * ( prob_2 + prob_2 + .....)
        cum_prob_2 = 1-poisson.cdf(inv-1,model.ARRIVAL_RATE * model.CUSTOMER_1_PROP)
        total_expected_revenue_in_period_2 += prob_1 * cum_prob_2 * expected_revenue_in_period_2_given_price(inv,p2,inv,0,model=model)

    return total_expected_revenue_in_period_2 + expected_revenue_in_period_1(model=model)

def get_best_dynamic_threshold(model:ModelInfo):
    best_rev = 0
    for inv in range(model.STARTING_INVENTORY+2):
        curr_rev = calculate_expected_total_revenue(c2_threshold=inv,model=model)
        if curr_rev >= best_rev:
            best_rev = curr_rev
            best_inv = inv
    return (best_inv,best_rev)

def get_static_pricing(model:ModelInfo):
    return (calculate_expected_total_revenue_given_low_low(model=model),calculate_expected_total_revenue_given_price(model.HIGH_PRICE,model=model),calculate_expected_total_revenue_given_price(model.LOW_PRICE,model=model))


def gap_between_dynamic_and_static(model:ModelInfo):
    tmp =max(get_static_pricing(model=model))
    return calculate_expected_total_revenue(get_best_dynamic_threshold(model=model)[0],model=model)-tmp


def High_Low_Gap(high_low,model:ModelInfo):
    res =  gap_between_dynamic_and_static(model=model.set_HIGH_PRICE(HIGH_PRICE=high_low[0]).set_LOW_PRICE(high_low[1]))
    return res


def Starting_Inv_Arrival_Gap(starting_inv_arrival,model:ModelInfo):
    res =  gap_between_dynamic_and_static(model=model.set_STARTING_INVENTORY(starting_inv_arrival[0]).set_ARRIVAL_RATE(starting_inv_arrival[1]))
    return res