#!/usr/bin/env python
# encoding: utf-8
"""Tests for stochastic cash flow models"""

import numpy
import matplotlib.pyplot as plt

from stochastic_prices import _load_prices, _simulate
from base_prices import PRICES, PRICECLASSES
from base_prices_2009 import PRICES_2009, PRICECLASSES_2009
from base_biomass_prices import BM_PRICES, BM_PRICECLASSES


YEARS = 100
ITERATIONS = 20
DPI = 100
FORMAT = 'svg'


def simulate_prices(I, N, proc, prices, priceclasses, BM, mean_price_col=None):
    """Simulate I price scenarios of N years with the given procedure proc
    (or equation)
    """
    meanprices = numpy.zeros((I, N), dtype=float)
    P_hat, delta_st = _load_prices(prices)
    steps = range(N)
    ptype = 'biomass' if BM else 'timber'
    for i in range(I):
        
        # simulate future prices using random process
        P_tilde = _simulate(P_hat, delta_st, N, proc, BM, priceclasses)
        
        # plot the simulated prices for the various assortments
        plt.figure(1, figsize=(8,4), dpi=DPI)
        for j in range(priceclasses):
            plt.plot(steps, P_tilde[:,j])
        ymin = numpy.min(P_tilde)
        ymax = numpy.max(P_tilde)
        plt.ylim((ymin, ymax,))
        plt.title("equation %s, iteration %d" % (proc, i+1))
        plt.ylabel('EUR')
        plt.xlabel('STEP')
        img_file = 'img/%s_eq%s_%d.%s' % (ptype, proc, i+1, FORMAT)
        plt.savefig(img_file, dpi=DPI)
        plt.close()
        
        # store the mean prices
        if mean_price_col is not None:
            mu = P_tilde[:, mean_price_col]
        else:
            mu = numpy.mean(P_tilde, axis=1)
        meanprices[i,:] = mu[:]
    
    # plot the mean prices into a single figure
    plt.figure(1, figsize=(8,5), dpi=DPI)
    ymin = numpy.min(meanprices[:,:])
    ymax = numpy.max(meanprices[:,:])
    for i in range(I):
        plt.plot(steps, meanprices[i,:])
    plt.ylim((ymin, ymax,))
    plt.title("equation %s" % (proc))
    plt.ylabel('EUR')
    plt.xlabel('STEP')
    img_file = 'img/%s_mean_eq%s.%s' % (ptype, proc, FORMAT)
    plt.savefig(img_file, dpi=DPI)
    plt.close()


def main():
    
    simulate_prices(ITERATIONS, YEARS, 2, PRICES, PRICECLASSES, False, -1)
    simulate_prices(ITERATIONS, YEARS, 5, PRICES, PRICECLASSES, False, -1)
    simulate_prices(ITERATIONS, YEARS, 8, PRICES, PRICECLASSES, False, -1)
    simulate_prices(ITERATIONS, YEARS, 'OU', PRICES_2009, PRICECLASSES_2009, False, None)
    
    #simulate_prices(ITERATIONS, YEARS, 2, BM_PRICES, BM_PRICECLASSES, True, -1)
    #simulate_prices(ITERATIONS, YEARS, 5, BM_PRICES, BM_PRICECLASSES, True, -1)
    #simulate_prices(ITERATIONS, YEARS, 8, BM_PRICES, BM_PRICECLASSES, True, -1)



if __name__ == '__main__':
    main()