"""Estimates the semilinear Choo and Siow homoskedastic (2006) model using Poisson GLM.
"""


import numpy as np
import scipy.linalg as spla
import scipy.sparse as spr
from sklearn import linear_model
from typing import Optional, Tuple
from dataclasses import dataclass

from entropy import _fill_hessianMuMu_from_components, _fill_hessianMuR_from_components
from choo_siow import entropy_choo_siow
from model_classes import ChooSiowPrimitives
from utils import print_stars
from matching_utils import Matching, _make_phi_mat,  _variance_muhat
from ipfp_solvers import ipfp_homoskedastic_solver


@dataclass
class PoissonGLMResults:
    """Stores and formats the estimation results.    """
    X: int
    Y: int
    K: int
    number_households: int
    number_individuals: int
    estimated_coefficients: np.ndarray
    estimated_Phi: np.ndarray
    estimated_u: np.ndarray
    estimated_v: np.ndarray
    variance_coefficients: np.ndarray
    stderrs_coefficients: np.ndarray

    def __repr__(self):
        line_stars = '*' * 80 + "\n"
        print_stars("Estimating a Choo and Siow model by Poisson GLM.")
        model_str = f"The data has {self.number_households} households\n\n"
        model_str += f"We use {self.K} basis functions.\n\n"
        repr_str = line_stars + model_str
        repr_str += "The estimated coefficients (and their standard errors) are\n\n"
        for i, coeff in enumerate(self.estimated_coefficients):
            repr_str += f"   base {i + 1}: {coeff: > 10.3f}  ({self.stderrs_coefficients[i]: .3f})\n"
        repr_str += "The estimated utilities of men are\n\n"
        for i, u_val in enumerate(self.estimated_u):
            repr_str += f"   base {i + 1}: {u_val: > 10.3f}\n"
        repr_str += "The estimated utilities of men are\n\n"
        for i, v_val in enumerate(self.estimated_v):
            repr_str += f"   base {i + 1}: {v_val: > 10.3f}\n"
        return repr_str + line_stars


def _prepare_data(muhat: Matching) -> Tuple[np.ndarray, int, int]:
    """Normalizes the matching patterns and stacks them.

       Args:
        muhat: the observed Matching
        phi_bases: an (X, Y, K) array of bases
        entropy: an ``EntropyFunctions`` object
        initial_omega: if specified, used as the weighting matrix for the first step when `param_entropy` is `True`

       Returns:
        the stacked muxy, mux0, mu0y
        the number of households
        the number of individuals

       Example:
        See the `main` function in the source of the module.
    """
    muxy, mux0, mu0y, n, m = muhat.unpack()
    n_couples = np.sum(muxy)
    n_households =  n_couples + np.sum(mux0) + np.sum(mu0y)
    n_individuals = n_households + n_couples
    # rescale the data so that the total number of individuals is one
    muhat_norm = np.concatenate(
        [muxy.flatten(), mux0, mu0y]) / n_individuals
    return muhat_norm, n_households, n_individuals


def choo_siow_poisson_glm(muhat: Matching, phi_bases: np.ndarray,
                          tol: Optional[float] = 1e-12,
                          max_iter: Optional[int] = 10000,
                          verbose: Optional[int] = 1):
    """
    Estimates the semilinear Choo and Siow homoskedastic (2006) model using Poisson GLM.

       Args:
        muhat: the observed Matching
        phi_bases: an (X, Y, K) array of bases
        tol: tolerance level for ``linear_model.PoissonRegressor.fit``
        max_iter: maximum number of iterations for ``linear_model.PoissonRegressor.fit``
        verbose: defines how much output we want (0 = least)

       Returns:
        a `PoissonGLMResults` instance

       Example:
        See the `main` function in the source of the module.
    """
    X, Y, K = phi_bases.shape

    XY = X * Y
    # the vector of weights for the Poisson regression
    d_a = np.append(2 * np.ones(XY), np.ones(X + Y))
    # the (XY+X+Y, X+Y) matrix
    B = spr.vstack([
        spr.hstack([
            spr.kron(spr.identity(X), np.ones((Y, 1))), spr.kron(
                np.ones((X, 1)), spr.identity(Y))
        ]),
        spr.hstack([
            spr.identity(X), spr.csr_matrix((X, Y))
        ]),
        spr.hstack([
            spr.csr_matrix((Y, X)), spr.identity(Y)
        ])
    ])
    # reshape the bases into an (XY+X+Y, K) matrix
    phi_mat = _make_phi_mat(phi_bases)
    phi_stacked = np.vstack(
        (phi_mat, np.zeros((X, K)), np.zeros((Y, K))))
    # finally, the (XY+X+Y, X+Y+K) matrix C
    C = spr.hstack([-B, phi_stacked])

    muxy, mux0, mu0y, n, m = muhat.unpack()
    muxyhat_norm, n_households, n_individuals = _prepare_data(muhat)

    clf = linear_model.PoissonRegressor(
        fit_intercept=False, tol=tol, verbose=verbose, alpha=0, max_iter=max_iter)
    clf.fit(spr.diags(1 / d_a) @ C, muxyhat_norm, sample_weight=d_a)

    beta_est = clf.coef_[-K:]

    # we compoute the variance-covariance of the estimator
    varmunm = _variance_muhat(muhat)
    hessian_components = entropy_choo_siow.e0_derivative(muhat)
    hessian_components_mumu, hessian_components_mur = hessian_components
    hessian_mumu = _fill_hessianMuMu_from_components(hessian_components_mumu)
    hessian_mur = _fill_hessianMuR_from_components(hessian_components_mur)
    inverse_mumu = spla.inv(hessian_mumu)
    phi_mat_T = phi_mat.T
    inv_mat = spla.inv(phi_mat_T @ inverse_mumu @ phi_mat)
    big_M_T = inv_mat @ np.concatenate((phi_mat_T, phi_mat_T @ inverse_mumu @ hessian_mur), axis=1)
    big_M = big_M_T.T
    variance_coefficients = big_M_T @ varmunm @ big_M
    stderrs_coefficients = np.sqrt(np.diag(variance_coefficients))

    Phi_est = phi_bases @ beta_est

    # we correct for the effect of the normalization
    u_est = clf.coef_[:X]  + np.log(n/n_individuals)
    v_est = clf.coef_[X: -K] + np.log(m/n_individuals)

    results = PoissonGLMResults(X=X, Y=Y, K=K, number_households=n_households, number_individuals=n_individuals,
                                estimated_coefficients=beta_est, estimated_Phi=Phi_est,
                                estimated_u=u_est, estimated_v=v_est,
                                variance_coefficients=variance_coefficients,
                                stderrs_coefficients=stderrs_coefficients)

    return results


if __name__ == '__main__':

    # Example 1: using IPFP to solve for the matching

    X, Y, K = 4, 3, 6
    phi_bases = np.zeros((X, Y, K))
    phi_bases[:, :, 0] = 1
    for x in range(X):
        phi_bases[x, :, 1] = x
        phi_bases[x, :, 3] = x*x
        for y in range(Y):
            phi_bases[x, y, 4] = x*y
    for y in range(Y):
        phi_bases[:, y, 2] = y
        phi_bases[:, y, 5] = y*y

    lambda_true = np.random.randn(K)
    phi_bases = np.random.randn(X, Y, K)
    Phi = phi_bases @ lambda_true

    nx = np.array([5, 6, 5, 8])*10000
    my = np.array([4, 8, 6])*10000
    (muxy, mux0, mu0y), err_x, err_y = ipfp_homoskedastic_solver(Phi, nx, my)
    results = choo_siow_poisson_glm(Matching(muxy, nx, my), phi_bases)
    print(results)

    estimates = results.estimated_coefficients
    stderrs = results.stderrs_coefficients

    repr_str = "The true and estimated coefficients (and their standard errors) are\n\n"
    for i, coeff in enumerate(estimates):
        repr_str += f"   base {i + 1}: {lambda_true[i]: > 10.3f}    {coeff: > 10.3f}  ({stderrs[i]: > 10.3f})\n"
    print_stars(repr_str)

    print_stars("The true and estimated utilities for men are:")
    print(np.column_stack((-np.log(mux0/nx), results.estimated_u)))

    print_stars("The true and estimated utilities for women are:")
    print(np.column_stack((-np.log(mu0y/my), results.estimated_v)))

    # Example 2: simulating many individuals
    n_households = 1e6
    # we simulate a Choo and Siow population with equal numbers of men and women of each type
    n = np.ones(X)
    m = np.ones(Y)
    Phi = phi_bases @ lambda_true
    choo_siow_instance = ChooSiowPrimitives(Phi, n, m)
    mus_sim  = choo_siow_instance.simulate(n_households)
    muxy_sim, mux0_sim, mu0y_sim, n_sim, m_sim = mus_sim.unpack()

    results = choo_siow_poisson_glm(mus_sim, phi_bases)
    print(results)

    estimates = results.estimated_coefficients
    stderrs = results.stderrs_coefficients

    repr_str = "The true and estimated coefficients (and their standard errors) are\n\n"
    for i, coeff in enumerate(estimates):
        repr_str += f"   base {i + 1}: {lambda_true[i]: > 10.3f}    {coeff: > 10.3f}  ({stderrs[i]: > 10.3f})\n"
    print_stars(repr_str)


    print_stars("The true and estimated utilities for men are:")
    print(np.column_stack((-np.log(mux0_sim/n_sim), results.estimated_u)))

    print_stars("The true and estimated utilities for women are:")
    print(np.column_stack((-np.log(mu0y_sim/m_sim), results.estimated_v)))