dns = np.loadtxt("data/DNSsol.dat")
utau = Retau*eqn.nu*2.0
ydns = dns[:,0]*0.5
udns = dns[:,2]*utau
f = interp1d(ydns, udns)
utarget = f(eqn.y)



beta_prior = eqn.beta.copy()
sigma_obs = 1e-10
sigma_prior = 1.0

eqn.objective = BayesianObjectiveU(utarget, beta_prior, sigma_obs, sigma_prior)
inverse_solver = InverseSolver(eqn)
inverse_solver.maxiter = 10
inverse_solver.stepsize = 0.1
inverse_solver.algo = "gn"
eqn = inverse_solver.solve()

plt.figure(11)
plt.semilogx(eqn.yp, up_prior, 'g-', label=r'Prior')
plt.semilogx(eqn.yp, eqn.up, 'r-', label=r'Posterior')
plt.semilogx(eqn.yp[::5], utarget[::5]/utau, 'b.', label=r'DNS')
plt.xlabel(r"$y^+$")
plt.ylabel(r"$u^+$")
plt.legend(loc=2)
plt.tight_layout()
plt.savefig("figs/inverse_matchu_u.pdf")
eqn.beta[:] = beta[:]
eqn.solve()

sigma_prior = 1.0
sigma_obs = 1e-10
q_target = eqn.q[:] + np.random.randn(np.size(eqn.q)) * sigma_obs

eqn.beta[:] = 2.0
eqn.solve()
q_prior = eqn.q.copy()
beta_prior = eqn.beta.copy()


eqn.objective = BayesianObjective(q_target, beta_prior, sigma_obs, sigma_prior)

inverse_solver = InverseSolver(eqn)
inverse_solver.maxiter = 100
eqn = inverse_solver.solve()

plt.figure(1)
plt.ioff()
plt.plot(y, q_prior, "g-", label="Prior")
plt.plot(y, eqn.q, "r-", label="Posterior")
plt.plot(y, q_target, "bx", label="Target")
plt.xlabel("y")
plt.ylabel(r"u")
plt.legend(loc=2)

plt.figure(2)
plt.plot(y, beta_prior, "g-", label="Prior")
plt.plot(y, eqn.beta, "r-", label="Posterior")
        eqn.beta = np.loadtxt("beta").astype(np.complex)
    else:
        eqn.beta[:] = 1.0

    beta_prior = np.ones_like(eqn.beta)
    sigma_obs = 1e-10
    sigma_prior = 1.0

    if objective[0] == "ALL":
        eqn.objective = BayesianObjectiveAll(qtarget, beta_prior, sigma_obs, sigma_prior)
    elif objective[0] == "U":
        eqn.objective = BayesianObjectiveU(qtarget, beta_prior, sigma_obs, sigma_prior)
    else:
        raise ValueError("Objective function not defined!")

    inverse_solver = InverseSolver(eqn)
    inverse_solver.maxiter = maxiter
    inverse_solver.stepsize = maxstep
    inverse_solver.algo = "sd"
    eqn = inverse_solver.solve()

    plt.ioff()
    plt.figure(11)
    plt.semilogx(eqn.yp, up_prior, "g-", label=r"Prior")
    plt.semilogx(eqn.yp, eqn.up, "r-", label=r"Posterior")
    plt.semilogx(dns.yp[::5], dns.u[::5], "b.", label=r"DNS")
    plt.xlabel(r"$y^+$")
    plt.ylabel(r"$u^+$")
    plt.legend(loc=2)
    plt.tight_layout()
    plt.savefig("figs/inverse_%s_u.pdf" % objective)
    if obs_cov_type == "scalar":
        objective.cov_obs_inv = np.linalg.inv(np.eye(ngrid-2)*sigma_obs**2)
    elif obs_cov_type == "vector":
        Cov_ = np.eye(ngrid-2)
        for i in range(ngrid-2):
            Cov_[i,i] = sigma_vector[i]**2
        objective.cov_obs_inv = np.linalg.inv(Cov_)
    elif obs_cov_type == "matrix":
        objective.cov_obs_inv = np.linalg.inv(Cov[1:-1,1:-1])
    else:
        raise ValueError("Wrong argument for covariance matrix type.")

    heat = HeatModel(T_inf, ngrid=ngrid)
    heat.solve()
    xi, Ti = heat.x, heat.T.copy()
    heat.objective = objective
    inverse = InverseSolver(heat)
    inverse.maxiter = 30
    inverse.nsamples = 2
    heat = inverse.solve()
    xf, Tf = heat.x, heat.T
    
    beta_map = heat.beta.astype(np.float64)
    beta_prior = beta_prior.astype(np.float64)
    data = data.astype(np.float64)
    sigma_obsv = np.ones_like(beta_map)*sigma_obs
    sigma_priorv = np.ones_like(beta_map)*sigma_prior
    np.savetxt("mcmc_input_cov.dat", Cov)
    np.savetxt('mcmc_input.dat', np.c_[data, np.sqrt(np.diag(Cov)), beta_prior, sigma_priorv, beta_map])

        elif obs_cov_type == "vector":
            Cov_ = np.eye(ngrid-2)
            for i in range(ngrid-2):
                Cov_[i,i] = sigma_vector[i]**2
                objective.cov_obs_inv = np.linalg.inv(Cov_)
        elif obs_cov_type == "matrix":
            objective.cov_obs_inv = np.linalg.inv(Cov[1:-1,1:-1])
        else:
            raise ValueError("Wrong argument for covariance matrix type.")
                
        heat = HeatModel(T_inf, ngrid=ngrid)
        heat.solve()
        xi, Ti = heat.x, heat.T.copy()
        
        heat.objective = objective
        inverse = InverseSolver(heat)
        inverse.nsamples = nsamples
        heat = inverse.solve()
        xf, Tf = heat.x, heat.T
        
        beta_prior_mean = np.mean(inverse.beta_samples_prior, axis=0)
        T_prior_mean = np.mean(inverse.T_samples_prior, axis=0)
        beta_prior_std = np.std(inverse.beta_samples_prior, axis=0)
        T_prior_std = np.std(inverse.T_samples_prior, axis=0)
        
        beta_post_mean = np.mean(inverse.beta_samples_post, axis=0)
        T_post_mean = np.mean(inverse.T_samples_post, axis=0)
        beta_post_std = np.std(inverse.beta_samples_post, axis=0)
        T_post_std = np.std(inverse.T_samples_post, axis=0)

f = interp1d(ydns, R11b)
qtarget[1::6] = f(eqn.y)
f = interp1d(ydns, R12b)
qtarget[2::6] = f(eqn.y)
f = interp1d(ydns, R22b)
qtarget[3::6] = f(eqn.y)
f = interp1d(ydns, R33b)
qtarget[4::6] = f(eqn.y)


eqn.beta = np.loadtxt("beta").astype(np.complex)
beta_prior = eqn.beta.copy()
sigma_obs = 1e-10
sigma_prior = 1.0
eqn.objective = BayesianObjectiveAll(qtarget, beta_prior, sigma_obs, sigma_prior)
inverse_solver = InverseSolver(eqn)
inverse_solver.maxiter = 10
inverse_solver.stepsize = 0.2
inverse_solver.algo = "sd"
eqn = inverse_solver.solve()

plt.figure(11)
plt.semilogx(eqn.yp, up_prior, 'g-', label=r'Prior')
plt.semilogx(eqn.yp, eqn.up, 'r-', label=r'Posterior')
plt.semilogx(dns.yp[::5], dns.u[::5], 'b.', label=r'DNS')
plt.xlabel(r"$y^+$")
plt.ylabel(r"$u^+$")
plt.legend(loc=2)
plt.tight_layout()
plt.savefig("figs/inverse_matchall_u.pdf")