import pyemu

"""
Explanation: Model background
Here is an example based on the Henry saltwater intrusion problem.  The synthetic model is a 2-dimensional SEAWAT model (X-Z domain) with 1 row, 120 columns and 20 layers.  The left boundary is a specified flux of freshwater, the right boundary is a specified head and concentration saltwater boundary.  The model has two stress periods: an initial steady state (calibration) period, then a transient period with less flux (forecast).  
The inverse problem has 603 parameters: 600 hydraulic conductivity pilot points, 1 global hydraulic conductivity, 1 specified flux multiplier for history matching and 1 specified flux multiplier for forecast conditions.  The inverse problem has 36 obseravtions (21 heads and 15 concentrations) measured at the end of the steady-state calibration period.  The forecasts of interest of the distance from the left model edge to the 10% seawater concentration in the basal model layer and the concentration at location 10.  Both of there forecasts are "measured" at the end of the forecast stress period.  The forecasts are both in the Jacobian matrix as zero-weight observations named pd_ten and C_obs10_2.I previously calculated the jacobian matrix, which is in the henry/ folder, along with the PEST control file.
Unlike the Schur's complement example notebook, here we will examine the consequences of not adjusting the specified flux multiplier parameters (mult1 and mult2) during inversion, since these types of model inputs are not typically considered for adjustment.
Using pyemu
End of explanation
"""


la = pyemu.ErrVar(jco=os.path.join("henry", "pest.jcb"),
                  omitted_parameters=["mult1","mult2"])
print(la.jco.shape) #without the omitted parameter or the prior info
la.forecast_names

"""
Explanation: First create a linear_analysis object.  We will use err_var  derived type, which replicates the behavior of the PREDVAR suite of PEST as well as ident_par utility.  We pass it the name of the jacobian matrix file.  Since we don't pass an explicit argument for parcov or obscov, pyemu attempts to build them from the parameter bounds and observation weights in a pest control file (.pst) with the same base case name as the jacobian.  Since we are interested in forecast uncertainty as well as parameter uncertainty, we also pass the names of the forecast sensitivity vectors we are interested in, which are stored in the jacobian as well.  Note that the forecasts argument can be a mixed list of observation names, other jacobian files or PEST-compatible ASCII matrix files.  Remember you can pass a filename to the verbose argument to write log file.
Since most groundwater model history-matching analyses focus on adjusting hetergeneous hydraulic properties and not boundary condition elements, let's identify the mult1 and mult2 parameters as omitted in the error variance analysis.  We can conceptually think of this action as excluding the mult1 and mult2 parameters from the history-matching process. Later we will explicitly calculate the penalty for not adjusting this parameter.
End of explanation
"""


s = la.qhalfx.s

import pylab as plt
figure = plt.figure(figsize=(10, 5))
ax = plt.subplot(111)
ax.plot(s.x)
ax.set_title("singular spectrum")
ax.set_ylabel("power")
ax.set_xlabel("singular value")
ax.set_xlim(0,20)
plt.show()

"""
Explanation: Parameter identifiability
The errvar dervied type exposes a method to get a pandas dataframe of parameter identifiability information.  Recall that parameter identifiability is expressed as $d_i = \Sigma(\mathbf{V}_{1i})^2$, where $d_i$ is the parameter identifiability, which ranges from 0 (not identified by the data) to 1 (full identified by the data), and $\mathbf{V}_1$ are the right singular vectors corresonding to non-(numerically) zero singular values.  First let's look at the singular spectrum of $\mathbf{Q}^{\frac{1}{2}}\mathbf{J}$, where $\mathbf{Q}$ is the cofactor matrix and $\mathbf{J}$ is the jacobian:
End of explanation
"""


ident_df = la.get_identifiability_dataframe(3) # the method is passed the number of singular vectors to include in V_1
ident_df.sort_values(by="ident").iloc[0:10]

"""
Explanation: We see that the singluar spectrum decays rapidly (not uncommon) and that we can really only support about 3 right singular vectors even though we have 600+ parameters in the inverse problem.  
Let's get the identifiability dataframe at 15 singular vectors:
End of explanation
"""


sing_vals = np.arange(40)

"""
Explanation: Plot the indentifiability:
We see that the global_k parameter has a much higher identifiability than any one of the 600 pilot points
Forecast error variance
Now let's explore the error variance of the forecasts we are interested in.  We will use an extended version of the forecast error variance equation:   
$\sigma_{s - \hat{s}}^2 = \underbrace{\textbf{y}i^T({\bf{I}} - {\textbf{R}})\boldsymbol{\Sigma}{{\boldsymbol{\theta}}i}({\textbf{I}} - {\textbf{R}})^T\textbf{y}_i}{1} + \underbrace{{\textbf{y}}i^T{\bf{G}}\boldsymbol{\Sigma}{\mathbf{\epsilon}}{\textbf{G}}^T{\textbf{y}}i}{2} + \underbrace{{\bf{p}}\boldsymbol{\Sigma}{{\boldsymbol{\theta}}_o}{\bf{p}}^T}{3}$
Where term 1 is the null-space contribution, term 2 is the solution space contribution and term 3 is the model error term (the penalty for not adjusting uncertain parameters).  Remember the mult1 and mult2 parameters that we marked as omitted?  The consequences of that action can now be explicitly evaluated.  See Moore and Doherty (2005) and White and other (2014) for more explanation of these terms.  Note that if you don't have any omitted_parameters, the only terms 1 and 2 contribute to error variance
First we need to create a list (or numpy ndarray) of the singular values we want to test.  Since we have $\lt40$ data, we only need to test up to $40$ singular values because that is where the action is:
End of explanation
"""


errvar_df = la.get_errvar_dataframe(sing_vals)
errvar_df.iloc[0:10]

"""
Explanation: The errvar derived type exposes a convience method to get a multi-index pandas dataframe with each of the terms of the error variance equation:
End of explanation
"""


fig = plt.figure(figsize=(10, 10))
ax_1, ax_2= plt.subplot(211), plt.subplot(212)
axes = [ax_1,ax_2]

colors = {"first": 'g', "second": 'b', "third": 'c'}
max_idx = 19
idx = sing_vals[:max_idx]
for ipred, pred in enumerate(la.forecast_names):
    pred = pred.lower()
    ax = axes[ipred]
    ax.set_title(pred)
    first = errvar_df[("first", pred)][:max_idx]
    second = errvar_df[("second", pred)][:max_idx]
    third = errvar_df[("third", pred)][:max_idx]
    ax.bar(idx, first, width=1.0, edgecolor="none", facecolor=colors["first"], label="first",bottom=0.0)
    ax.bar(idx, second, width=1.0, edgecolor="none", facecolor=colors["second"], label="second", bottom=first)
    ax.bar(idx, third, width=1.0, edgecolor="none", facecolor=colors["third"], label="third", bottom=second+first)
    ax.set_xlim(-1,max_idx+1)
    ax.set_xticks(idx+0.5)
    ax.set_xticklabels(idx)
    if ipred == 2:
        ax.set_xlabel("singular value")
    
    ax.set_ylabel("error variance")
    ax.legend(loc="upper right")
plt.show()


"""
Explanation: plot the error variance components for each forecast:
End of explanation
"""


schur = la.get(astype=pyemu.Schur)
schur_prior = schur.prior_forecast
schur_post = schur.posterior_forecast
print("{0:10s} {1:>12s} {2:>12s} {3:>12s} {4:>12s}"
      .format("forecast","errvar prior","errvar min",
              "schur prior", "schur post"))
for ipred, pred in enumerate(la.forecast_names):
    first = errvar_df[("first", pred)][:max_idx]
    second = errvar_df[("second", pred)][:max_idx]  
    min_ev = np.min(first + second)
    prior_ev = first[0] + second[0]
    prior_sh = schur_prior[pred]
    post_sh = schur_post[pred]
    print("{0:12s} {1:12.6f} {2:12.6f} {3:12.6} {4:12.6f}"
          .format(pred,prior_ev,min_ev,prior_sh,post_sh))

"""
Explanation: Here we see the trade off between getting a good fit to push down the null-space (1st) term and the penalty for overfitting (the rise of the solution space (2nd) term)).  The sum of the first two terms in the "appearent" error variance (e.g. the uncertainty that standard analyses would yield) without considering the contribution from the omitted parameters.  You can verify this be checking prior uncertainty from the Schur's complement notebook against the zero singular value result using only terms 1 and 2.
We also see the added penalty for not adjusting the mult1 and mult2 parameters (3rd term).  The ability to forecast the distance from the left edge of the model to the 10% saltwater concentration and the forecast the concentration at location 10 has been compromised by not adjusting mult1 and mult2 during calibration.  
Let's check the errvar results against the results from schur. This is simple with pyemu, we simply  cast the errvar type to a schur type:
End of explanation
"""


