"""
Filename: odu.py

Authors: Thomas Sargent, John Stachurski

Solves the "Offer Distribution Unknown" Model by value function
iteration and a second faster method discussed in the corresponding
quantecon lecture.

"""
from textwrap import dedent
from scipy.interpolate import LinearNDInterpolator
from scipy.integrate import fixed_quad
from scipy.stats import beta as beta_distribution
from scipy import interp
from numpy import maximum as npmax
import numpy as np


class SearchProblem(object):
    """
    A class to store a given parameterization of the "offer distribution
    unknown" model.

    Parameters
    ----------
    beta : scalar(float), optional(default=0.95)
        The discount parameter
    c : scalar(float), optional(default=0.6)
        The unemployment compensation
    F_a : scalar(float), optional(default=1)
        First parameter of beta distribution on F
    F_b : scalar(float), optional(default=1)
        Second parameter of beta distribution on F
    G_a : scalar(float), optional(default=3)
        First parameter of beta distribution on G
    G_b : scalar(float), optional(default=1.2)
        Second parameter of beta distribution on G
    w_max : scalar(float), optional(default=2)
        Maximum wage possible
    w_grid_size : scalar(int), optional(default=40)
        Size of the grid on wages
    pi_grid_size : scalar(int), optional(default=40)
        Size of the grid on probabilities

    Attributes
    ----------
    beta, c, w_max : see Parameters
    w_grid : np.ndarray
        Grid points over wages, ndim=1
    pi_grid : np.ndarray
        Grid points over pi, ndim=1
    grid_points : np.ndarray
        Combined grid points, ndim=2
    F : scipy.stats._distn_infrastructure.rv_frozen
        Beta distribution with params (F_a, F_b), scaled by w_max
    G : scipy.stats._distn_infrastructure.rv_frozen
        Beta distribution with params (G_a, G_b), scaled by w_max
    f : function
        Density of F
    g : function
        Density of G
    pi_min : scalar(float)
        Minimum of grid over pi
    pi_max : scalar(float)
        Maximum of grid over pi
    """

    def __init__(self, beta=0.95, c=0.6, F_a=1, F_b=1, G_a=3, G_b=1.2,
                 w_max=2, w_grid_size=40, pi_grid_size=40):

        self.beta, self.c, self.w_max = beta, c, w_max
        self.F = beta_distribution(F_a, F_b, scale=w_max)
        self.G = beta_distribution(G_a, G_b, scale=w_max)
        self.f, self.g = self.F.pdf, self.G.pdf    # Density functions