import pymc3 as pm
import theano.tensor as tt

Ks=np.array([8.0,2.15])

with pm.Model() as model:

    # Common "shared" orbital parameters
    logK = pm.Uniform("logK",lower=np.log(0.5), upper=np.log(100.0), testval=np.log(Ks), shape=2)

    # Mass
#        logmpl = pm.Normal("logmpl", mu=np.log(float(msini1.value)*M_star[0],float(msini2.value)*M_star[0]), sd=(5.0,5.0),shape=2)
#        mpl = pm.Deterministic("mpl", tt.exp(logmpl))

    BoundedNormal = pm.Bound(pm.Normal, lower=1.0, upper=2000.0)
    P = BoundedNormal("P", mu=np.array(periods), sd=np.array(period_errs), shape=2,
                      testval=np.array(periods))

    #ecc = pm.Uniform("ecc", lower=0, upper=1, testval=0.1,shape=2)
    ecc = pm.Beta("ecc", alpha=0.867, beta=3.03, shape=2, testval=np.array([0.01, 0.15]))
    omega = xo.distributions.Angle("omega",shape=2)
    t0 = pm.Normal("t0", mu=np.array(t0s), sd=np.array(t0_errs), shape=2)

    # Track some parameters
    K = pm.Deterministic("K", pm.math.exp(logK))

    # We need to include a linear trend too
    trend = pm.Normal("trend", mu=0.0, sd=100.0)

    # Set up the orbit
    orbit = xo.orbits.KeplerianOrbit(period=P, t0=t0, ecc=ecc, omega=omega)

   # Loop over datasets
    logss = []
    sys_vels = []
    gp_params = []
    for t in utel:
        xt = x[tel== t]
        yt = y[tel == t]
        yerrt = yerr[tel == t]

        # Jitter for this instrument
        logs = pm.Uniform("logs_{0}".format(t), lower=-15, upper=15, testval=0.0) #<--- CHANGE?
        logss.append(logs)

        # Systemic velocity for this instrument
        sys_vel = pm.Uniform("sys_vel_{0}".format(t),lower=-100,upper=100,testval=0.0)
        sys_vels.append(sys_vel)

        # Compute the model RV
        vrad = orbit.get_radial_velocity(xt, K=K)
        #pm.Deterministic("vrad", vrad)

        # And the background model
        bkg_rv = sys_vel + trend * (xt - x_ref)
        #bkg = pm.Deterministic("bkg_rv",bkg_rv)

        # Sum over planets and add the background to get the full model
        #rv_model = pm.Deterministic("rv_model", tt.sum(vrad, axis=-1) + bkg_rv)

        # Noise model (error bar + jitter)
        err = pm.math.sqrt(yerrt**2 + pm.math.exp(2*logs))
        mod = tt.sum(vrad, axis=-1) + bkg_rv

        #GP
        logS1 = pm.Normal("logS1_{0}".format(t), mu=0.0, sd=15.0, testval=np.log(np.var(yt)))
        logw1 = pm.Normal("logw1_{0}".format(t), mu=0.0, sd=15.0, testval=np.log(3.0))
        logS2 = pm.Normal("logS2_{0}".format(t), mu=0.0, sd=15.0, testval=np.log(np.var(yt)))
        logw2 = pm.Normal("logw2_{0}".format(t), mu=0.0, sd=15.0, testval=np.log(3.0))
        logQ = pm.Normal("logQ_{0}".format(t), mu=0.0, sd=15.0, testval=0)
        gp_params += [logS1, logw1, logS2, logw2, logQ]

        # Set up the kernel an GP
        kernel = xo.gp.terms.SHOTerm(log_S0=logS1, log_w0=logw1, Q=1.0/np.sqrt(2))
        kernel += xo.gp.terms.SHOTerm(log_S0=logS2, log_w0=logw2, log_Q=logQ)
        gp = xo.gp.GP(kernel, xt, err)

        # Add a custom "potential" (log probability function) with the GP likelihood
        pm.Potential("rv_obs_{0}".format(t), gp.log_likelihood(yt-mod))
        pm.Deterministic("gp_pred_{0}".format(t), gp.predict())


    # Do 'pfs2' explicitly so can try to connect pfs1 and pfs2 sys_vel terms
  #  xt = x[tel== 'pfs2']
  #  yt = y[tel == 'pfs2']
  #  yerrt = yerr[tel == 'pfs2']

  #  logs = pm.Uniform("logs_{0}".format('pfs2'), lower=-5.0, upper=5.0, testval=0.0)
    #logs = pm.Normal("logs_{0}".format('pfs2'), mu=np.log(10.0), sd=10.0)
  #  logss.append(logs)

   # BoundedNormal = pm.Bound(pm.Normal, lower=sys_vel-2.0, upper=sys_vel+2.0)
   # sys_vel = BoundedNormal("sys_vel_{0}".format('pfs2'),mu=sys_vel, sd=2.0)
   # sys_vels.append(sys_vel)

  #  vrad = orbit.get_radial_velocity(xt, K=K)

  #  bkg_rv = sys_vel + trend * (xt - x_ref)

  #  err = pm.math.sqrt(yerrt**2 + pm.math.exp(2*logs))
  #  pm.Normal("obs_{0}".format('pfs2'), mu=tt.sum(vrad, axis=-1) + bkg_rv, sd=err, observed=yt)

    pm.Normal("pfs", mu=0, sd=3.0, observed=model.sys_vel_pfs1 - model.sys_vel_pfs2)
    pm.Deterministic("rv_model", orbit.get_radial_velocity(x, K=K))


    # These are for plotting
    rv_model_pred = []
    for i in range(2):
        x_pred = P[i] * phase_pred + t0[i]
        rv_model_pred.append(orbit.get_radial_velocity(x_pred, K=K)[:, i])
    pm.Deterministic("rv_model_pred", tt.stack(rv_model_pred))
    pm.Deterministic("gp_pred", gp.predict())


    print("\nOptimizing GP params...")
    map_soln = xo.optimize(vars=sys_vels + logss)
    map_soln = xo.optimize(map_soln, vars=[t0])
    map_soln = xo.optimize(map_soln, vars=gp_params)
    map_soln = xo.optimize(map_soln, vars=[t0, P, logK, ecc, omega])
    map_soln = xo.optimize(map_soln)