state_noise_dist = scipy.stats.multivariate_normal(cov=params.Q)
meas_noise_dist = scipy.stats.multivariate_normal(cov=params.R2)

params.xs[:, 0] = init_state
params.xsnofix[:, 0] = init_state

params.ys2[:, 0] = params.C2 @ params.xs[:, 0] + meas_noise_dist.rvs(
)  # measured from actual plant
SPF.init_filter(particles, params.ys2[:, 0], cstr_filter)

for k in range(numSwitches):
    switchtrack[k,
                0] = numpy.sum(particles.w[numpy.where(particles.s == k)[0]])

maxtrack[:, 0] = SPF.get_max_track(particles, numSwitches)
smoothedtrack[:, 0] = RBPF.smoothed_track(numSwitches, switchtrack, 1, 10)

params.spfmeans[:, 0], params.spfcovars[:, :, 0] = SPF.get_stats(particles)

# Loop through the rest of time
for t in range(1, params.N):
    random_element = state_noise_dist.rvs()
    if params.ts[t] < 50.0:
        params.xs[:, t] = params.cstr_model.run_reactor(
            params.xs[:, t - 1], params.us[t - 1], params.h) + random_element
        params.xsnofix[:, t] = params.cstr_model.run_reactor(
            params.xsnofix[:, t - 1], params.us[t - 1], params.h)
        params.xsnofix[:, t] += random_element  # actual plant
    else:
        params.xs[:, t] = params.cstr_model_broken.run_reactor(
            params.xs[:, t - 1], params.us[t - 1], params.h)
import typing

tend = 200
params = params.Params(tend)

# Divide state space into sectors: n by m
# Divide state space into sectors: n by m
nX = 2  # rows
nY = 2  # cols
xspace = numpy.array([0.0, 1.0])
yspace = numpy.array([250, 550])

linsystems = params.cstr_model.get_linear_systems(nX, nY, xspace, yspace,
                                                  params.h)

models, _ = RBPF.setup_rbpf(linsystems, params.C2, params.Q, params.R2)
A = numpy.array([[0.98, 0.00, 0.00, 0.00, 0.00, 0.01, 0.00],
                 [0.00, 0.98, 0.00, 0.01, 0.01, 0.01, 0.00],
                 [0.01, 0.00, 0.98, 0.00, 0.00, 0.01, 0.01],
                 [0.00, 0.00, 0.00, 0.98, 0.00, 0.01, 0.00],
                 [0.00, 0.01, 0.00, 0.00, 0.99, 0.00, 0.00],
                 [0.01, 0.01, 0.01, 0.01, 0.00, 0.96, 0.00],
                 [0.00, 0.00, 0.01, 0.00, 0.00, 0.00, 0.99]])
numModels = len(models)

nP = 500  # number of particles
init_state = linsystems[5].op

sguess = RBPF.get_initial_switches(init_state, linsystems)
particles = RBPF.init_rbpf(sguess, init_state, params.init_state_covar, 2, nP)
def fun():
    isDone = True
    init_state = numpy.array([0.55, 450])  # initial state

    # Setup Switching Particle Filter
    A = numpy.array([[0.999, 0.001], [0.001, 0.999]])

    def fun1(x, u, w):
        return params.cstr_model.run_reactor(x, u, params.h) + w

    def fun2(x, u, w):
        return params.cstr_model_broken.run_reactor(x, u, params.h) + w

    def gs(x):
        return params.C2 @ x

    F = [fun1, fun2]
    G = [gs, gs]
    numSwitches = 2

    ydists = numpy.array([
        scipy.stats.multivariate_normal(cov=params.R2),
        scipy.stats.multivariate_normal(cov=params.R2)
    ])
    xdists = numpy.array([
        scipy.stats.multivariate_normal(cov=params.Q),
        scipy.stats.multivariate_normal(cov=params.Q)
    ])
    cstr_filter = SPF.Model(F, G, A, xdists, ydists)

    nP = 500  # number of particles
    xdist = scipy.stats.multivariate_normal(mean=init_state,
                                            cov=params.init_state_covar)
    sdist = [0.9, 0.1]
    particles = SPF.init_spf(xdist, sdist, nP, 2)

    switchtrack = numpy.zeros([2, params.N])
    maxtrack = numpy.zeros([numSwitches, params.N])
    smoothedtrack = numpy.zeros([numSwitches, params.N])

    state_noise_dist = scipy.stats.multivariate_normal(cov=params.Q)
    meas_noise_dist = scipy.stats.multivariate_normal(cov=params.R2)

    # Setup control (use linear control)
    linsystems = params.cstr_model.get_nominal_linear_systems(params.h)
    linsystems_broken = params.cstr_model_broken.get_nominal_linear_systems(
        params.h)
    opoint = 1  # the specific linear model we will use

    lin_models = [None] * 2  # type: typing.List[RBPF.Model]
    lin_models[0] = RBPF.Model(linsystems[opoint].A, linsystems[opoint].B,
                               linsystems[opoint].b, params.C2, params.Q,
                               params.R2)
    lin_models[1] = RBPF.Model(linsystems_broken[opoint].A,
                               linsystems_broken[opoint].B,
                               linsystems_broken[opoint].b, params.C2,
                               params.Q, params.R2)

    H = numpy.matrix([1, 0])  # only attempt to control the concentration
    setpoint = 0.49
    controllers = [None] * 2  # type: typing.List[LQR.Controller]
    for k in range(2):
        ysp = setpoint - lin_models[k].b[0]  # set point is set here
        x_off, u_off = LQR.offset(lin_models[k].A,
                                  numpy.matrix(lin_models[k].B), params.C2, H,
                                  numpy.matrix([ysp]))
        K = LQR.lqr(lin_models[k].A, numpy.matrix(lin_models[k].B), params.QQ,
                    params.RR)
        controllers[k] = LQR.Controller(K, x_off, u_off)

    # Setup simulation
    params.xs[:, 0] = init_state
    params.ys2[:, 0] = params.C2 @ params.xs[:, 0] + meas_noise_dist.rvs(
    )  # measured from actual plant

    SPF.init_filter(particles, params.ys2[:, 0], cstr_filter)

    for k in range(numSwitches):
        switchtrack[k, 0] = numpy.sum(
            particles.w[numpy.where(particles.s == k)[0]])

    maxtrack[:, 0] = SPF.get_max_track(particles, numSwitches)
    smoothedtrack[:, 0] = RBPF.smoothed_track(numSwitches, switchtrack, 1, 10)

    params.spfmeans[:, 0], params.spfcovars[:, :, 0] = SPF.get_stats(particles)

    # Controller Input
    ind = numpy.argmax(maxtrack[:, 0])  # use this model and controller
    horizon = 150
    params.us[0] = MPC.mpc_lqr(params.spfmeans[:, 0] - lin_models[ind].b,
                               horizon, lin_models[ind].A,
                               numpy.matrix(lin_models[ind].B), params.QQ,
                               params.RR, controllers[ind].x_off,
                               controllers[ind].u_off)

    # Loop through the rest of time
    for t in range(1, params.N):
        random_element = state_noise_dist.rvs()
        if params.ts[t] < 100:  # break here
            params.xs[:, t] = params.cstr_model.run_reactor(
                params.xs[:, t - 1], params.us[t - 1],
                params.h) + random_element
        else:
            params.xs[:, t] = params.cstr_model_broken.run_reactor(
                params.xs[:, t - 1], params.us[t - 1], params.h)
            params.xs[:, t] += random_element

        params.ys2[:, t] = params.C2 @ params.xs[:, t] + meas_noise_dist.rvs(
        )  # measured from actual plant

        SPF.spf_filter(particles, params.us[t - 1], params.ys2[:, t],
                       cstr_filter)
        params.spfmeans[:, t], params.spfcovars[:, :,
                                                t] = SPF.get_stats(particles)

        for k in range(numSwitches):
            switchtrack[k, 0] = numpy.sum(
                particles.w[numpy.where(particles.s == k)[0]])

        maxtrack[:, t] = SPF.get_max_track(particles, numSwitches)
        smoothedtrack[:, t] = RBPF.smoothed_track(numSwitches, switchtrack, t,
                                                  40)

        # Controller Input
        if t % 1 == 0:
            ind = numpy.argmax(maxtrack[:, t])  # use this model and controller
            params.us[t] = MPC.mpc_lqr(
                params.spfmeans[:, t] - lin_models[ind].b, horizon,
                lin_models[ind].A, numpy.matrix(lin_models[ind].B), params.QQ,
                params.RR, controllers[ind].x_off, controllers[ind].u_off)
        else:
            params.us[t] = params.us[t - 1]
            if params.us[t] is None or numpy.isnan(params.us[t]):
                isDone = False
                break
    return isDone, setpoint
import numpy
import scipy.stats
import openloop.params as params
import src.Results as Results
import src.RBPF as RBPF
import matplotlib.pyplot as plt

tend = 150
params = params.Params(tend)

init_state = numpy.array([0.5, 400])  # initial state

linsystems = params.cstr_model.get_nominal_linear_systems(params.h)

models, _ = RBPF.setup_rbpf(linsystems, params.C1, params.Q, params.R1)
A = numpy.array([[0.5, 0.25, 0.25],
                 [0.25, 0.5, 0.25],
                 [0.25, 0.25, 0.5]])
numModels = len(models)

nP = 500  # number of particles
sguess = RBPF.get_initial_switches(init_state, linsystems)
particles = RBPF.init_rbpf(sguess, init_state, params.init_state_covar, 2, nP)

switchtrack = numpy.zeros([len(linsystems), params.N])
maxtrack = numpy.zeros([len(linsystems), params.N])
smoothedtrack = numpy.zeros([len(linsystems), params.N])

state_noise_dist = scipy.stats.multivariate_normal(cov=params.Q)
meas_noise_dist = scipy.stats.multivariate_normal(cov=params.R1)
switchtrack = numpy.zeros([2, params.N])
maxtrack = numpy.zeros([numSwitches, params.N])
smoothedtrack = numpy.zeros([numSwitches, params.N])

state_noise_dist = scipy.stats.multivariate_normal(cov=params.Q)
meas_noise_dist = scipy.stats.multivariate_normal(cov=params.R2)

# Setup control (use linear control)
linsystems = params.cstr_model.get_nominal_linear_systems(params.h)
linsystems_broken = params.cstr_model_broken.get_nominal_linear_systems(
    params.h)
opoint = 1  # the specific linear model we will use

lin_models = [None] * 2  # type: typing.List[RBPF.Model]
lin_models[0] = RBPF.Model(linsystems[opoint].A, linsystems[opoint].B,
                           linsystems[opoint].b, params.C2, params.Q,
                           params.R2)
lin_models[1] = RBPF.Model(linsystems_broken[opoint].A,
                           linsystems_broken[opoint].B,
                           linsystems_broken[opoint].b, params.C2, params.Q,
                           params.R2)

setpoint = linsystems[opoint].op
H = numpy.matrix([1, 0])  # only attempt to control the concentration
usps = numpy.zeros([len(lin_models), 1])
for k in range(len(lin_models)):
    sp = setpoint[0] - lin_models[k].b[0]
    x_off, usp = LQR.offset(lin_models[k].A, numpy.matrix(lin_models[k].B),
                            params.C2, H, numpy.matrix([sp]))
    ysp = x_off
    usp = numpy.array([usp])