module PharmValuation
# import mean and disallow method extension 
using Statistics: mean
using Printf: @printf
using Plots

struct CashProcess{T}
    annual_cash_flow::T
    drift::T 
    volatility::T 
    terminal_multiplier::T 
    risk_premium::T 
end

struct CostProcess{T}
    investment::T 
    total_cost::T 
    volatility::T 
    failure_rate::T 
end

struct SimOptions{T,U}
    time_step::T 
    term::U 
    runs::U 
    basis::Function 
end

struct ProjectProcess{T}
    CashProcess
    CostProcess
    correlation::T 
    risk_free_rate::T 
    SimOptions 
end

function main(;runs=100_000) 
    basis(x,y) = [1, x, y, x*y, x^2, y^2, x^2 * y, x * y^2, (x*y)^2]
    cash = CashProcess(20e6,0.02,35/100,5.0,0.036)
    cost = CostProcess(10e6,100e6,0.5,0.06931)
    sim = SimOptions(0.25,20,runs,basis)
    project = ProjectProcess(cash,cost,-0.1,0.05,sim)
    project_value = lsm(project)
    @printf "The project value: \$%.0f" project_value 
end

""" simulate returns the cash flow and cost matrices generated by their respective processes """
function simulate(pp::ProjectProcess)
    # Set the number of periods 
    periods = Int(ceil(pp.SimOptions.term / pp.SimOptions.time_step))

    # Risk adjusted cash flow drift rate 
    adj_cash_flow = pp.CashProcess.drift - pp.CashProcess.risk_premium

    # Set return matrices and correlated random errors.
    net_cash = randn(pp.SimOptions.runs, periods)
    cost = pp.correlation * net_cash .+ sqrt(1 - pp.correlation^2) * randn(pp.SimOptions.runs,periods)

    # simulate the investment costs and correlated cash flows 
    @inbounds for run in 1:pp.SimOptions.runs
        @inbounds for period in 1:periods
            # cash flow simulation
            prev_cash = period != 1 ? net_cash[run,period-1] : pp.CashProcess.annual_cash_flow
            net_cash[run,period] = prev_cash * exp((adj_cash_flow-0.5 * pp.CashProcess.volatility^2)*pp.SimOptions.time_step +
                pp.CashProcess.volatility * sqrt(pp.SimOptions.time_step) * net_cash[run,period])
            
            # cost simulation 
            prev_cost = period != 1 ? cost[run,period-1] : pp.CostProcess.total_cost
            current_cost = 0.0
            if prev_cost != 0
                current_cost = prev_cost - pp.CostProcess.investment * pp.SimOptions.time_step +
                    pp.CostProcess.volatility * sqrt(pp.CostProcess.investment*prev_cost*pp.SimOptions.time_step) * cost[run,period]
                current_cost = current_cost < 0 ? 0 : current_cost
            end
            cost[run,period] = current_cost 
        end
    end
    return net_cash, cost 
end

""" lsm returns the least squares monte carlo estimate for the optimal stopping problem"""
function lsm(pp::ProjectProcess)
    periods = Int(ceil(pp.SimOptions.term / pp.SimOptions.time_step))

    # local for clarity
    period_investmt = pp.CostProcess.investment * pp.SimOptions.time_step

    # Simulate the net cash flows and cost values 
    cash_matrix, cost_matrix = simulate(pp)

    # Set value function array to terminal value where investment has stopped. 
    value_func = zeros(size(cash_matrix))
    terminal_cash_flow = cash_matrix[:,end]
    value_func[:,end] = pp.CashProcess.terminal_multiplier * (cost_matrix[:,end] .== 0) .* terminal_cash_flow

    # Discount rates for investing phase and sales phase 
    cash_discount_rate = exp(-1 * pp.risk_free_rate * pp.SimOptions.time_step)
    invest_discount_rate = exp(-1 * (pp.risk_free_rate + pp.CostProcess.failure_rate)*pp.SimOptions.time_step)

    # value function iteration 
    @inbounds for period in (periods-1):-1:1
        # initialize next periods value 
        next_val = invest_discount_rate * value_func[:,period+1]

        # Apply basis function to cost and cash flows to approximate the discounted 
        # next period value. 
        basis_vecs = pp.SimOptions.basis.(cash_matrix[:,period],cost_matrix[:,period])

        # Transform vector of vector bases into a basis matrix for regression analysis.
        basis_matrix = reduce(hcat,basis_vecs)' 

        # Solve regression coefficients
        coefficients = basis_matrix \ next_val

        # Estimate continuation value of investment 
        est_val = basis_matrix * coefficients
        
        # Update value function array 
        @inbounds for run in 1:pp.SimOptions.runs
            # Investing phase 
            if cost_matrix[run,period] != 0
                # Investment decision criterion
                invest_val = est_val[run] - period_investmt
                
                # Update value function with for the investment phase.
                value_func[run,period] = invest_val > 0 ? next_val[run] - period_investmt : 0.0
            else
                # Update value function for the post-investment sales phase.
                value_func[run, period] = cash_matrix[run,period] * pp.SimOptions.time_step + 
                    cash_discount_rate * value_func[run,period+1]
            end
        end
    end
    retval = invest_discount_rate * value_func[:,1]
    return mean(retval)
end

end # module