import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

# function returns dsbdt, debdt, dibdt, dqbdt, drbdt, dszdt, dezdt, dizdt, dqzdt, drzdt
def i_seiqr(x, t):
    ab = 1                  # the number of new computers added to the network at each moment in time in Blue
    az = 1                  # the number of new computers added to the network at each moment in time in Red
    alfab = 0.1             # the rate of removal of a computer from the network owing to the action of the malware in Blue
    alfaz = 0.1             # the rate of removal of a computer from the network owing to the action of the malware in Red
    db = 0.1                # the rate of removal of computers from the network owing to causes not due to the malware in Blue
    dz = 0.1                # the rate of removal of computers from the network owing to causes not due to the malware in Red
    etab = 0.2              # the rate of the loss of immunity in Blue
    etaz = 0.2              # the rate of the loss of immunity in Red
    cb = 0.3                # the rate of passage from the susceptible state to the infected stage in Blue
    cz = 0.3                # the rate of passage from the susceptible state to the infected stage in Red
    deltab = 3.8            # the rate of passage from the infected state to the quarantine stage in Blue
    deltaz = 3.8            # the rate of passage from the infected state to the quarantine stage in Red
    mub = 0.3               # the rate of passage from the exposed state to the infected stage in Blue
    muz = 0.3               # the rate of passage from the exposed state to the infected stage in Red
    ksib = 0.3              # the rate of passage from the infectious or quarantined state to the recovered state in Blue
    ksiz = 0.3              # the rate of passage from the infectious or quarantined state to the recovered state in Red
    gamab = 1.8             # the rate of recovery due to the action of the antivirus software in Blue
    gamaz = 1.8             # the rate of recovery due to the action of the antivirus software in Red

    rhoup = 0.1             # Normal kinetic attack rate of Red forces
    rhodown = 0.01          # Diminished kinetic attack rate (because of infection) of Red forces
    bup = 0.1               # Normal kinetic attack rate of Blue forces
    bdown = 0.01            # Diminished kinetic attack rate (because of infection) of Blue forces

    sb = x[0]               # Initial number of Susceptible in Blue
    eb = x[1]               # Initial number of Exposed in Blue
    ib = x[2]               # Initial number of Infected in Blue
    qb = x[3]               # Initial number of Quarantined in Blue
    rb = x[4]               # Initial number of Recovered in Blue
    sz = x[5]               # Initial number of Susceptible in Red
    ez = x[6]               # Initial number of Exposed in Red
    iz = x[7]               # Initial number of Infected in Red
    qz = x[8]               # Initial number of Quarantined in Red
    rz = x[9]               # Initial number of Recovered in Red

    dsbdt = ab - cb * sb * ib - db * sb + etab * rb - (rhoup * (sz + ez + rz) + rhodown * (qz + iz)) * (sb / (sb + eb + ib + qb + rb))
    debdt = cb * sb * ib - (db + mub) * eb - (rhoup * (sz + ez + rz) + rhodown * (qz + iz)) * (eb / (sb + eb + ib + qb + rb))
    dibdt = mub * eb - (db + alfab + gamab + deltab) * ib - (rhoup * (sz + ez + rz) + rhodown * (qz + iz)) * (ib / (sb + eb + ib + qb + rb))
    dqbdt = deltab * ib - (db + alfab + ksib) * qb - (rhoup * (sz + ez + rz) + rhodown * (qz + iz)) * (qb / (sb + eb + ib + qb + rb))
    drbdt = gamab * ib + ksib * qb - (db + etab) * rb - (rhoup * (sz + ez + rz) + rhodown * (qz + iz)) * (rb / (sb + eb + ib + qb + rb))

    dszdt = az - cz * sz * iz - dz * sz + etaz * rz - (bup * (sb + eb + rb) + bdown * (qb + ib)) * (sz / (sz + ez + iz + qz + rz))
    dezdt = cz * sz * iz - (dz + muz) * ez - (bup * (sb + eb + rb) + bdown * (qb + ib)) * (ez / (sz + ez + iz + qz + rz))
    dizdt = muz * ez - (dz + alfaz + gamaz + deltaz) * iz - (bup * (sb + eb + rb) + bdown * (qb + ib)) * (iz / (sz + ez + iz + qz + rz))
    dqzdt = deltaz * iz - (dz + alfaz + ksiz) * qz - (bup * (sb + eb + rb) + bdown * (qb + ib)) * (qz / (sz + ez + iz + qz + rz))
    drzdt = gamaz * iz + ksiz * qz - (dz + etaz) * rz - (bup * (sb + eb + rb) + bdown * (qb + ib)) * (rz / (sz + ez + iz + qz + rz))

    return [dsbdt, debdt, dibdt, dqbdt, drbdt, dszdt, dezdt, dizdt, dqzdt, drzdt]


# initial condition
# Blue has 190 Susceptible, 0 Exposed/Infected/Quarantined and 10 Recovered.
# Red has 150 Susceptible, 10 Exposed, 20 Infected, 20 Quarantined and 0 Recovered.
x0 = [190, 0, 0, 0, 10, 150, 10, 20, 20, 0]
t = np.linspace(0, 20, 100)

# solve ODE
x = odeint(i_seiqr, x0, t)

sb = x[:, 0]
eb = x[:, 1]
ib = x[:, 2]
qb = x[:, 3]
rb = x[:, 4]
sz = x[:, 5]
ez = x[:, 6]
iz = x[:, 7]
qz = x[:, 8]
rz = x[:, 9]

# plot results
# for blue units
# plt.plot(t, sb, 'g-')
# plt.plot(t, eb, 'c--')
# plt.plot(t, ib, 'm-.')
# plt.plot(t, qb, 'y.')
# plt.plot(t, rb, 'k:')
plt.plot(t, sb + eb + ib + qb + rb, 'b-')

# for red units
# plt.plot(t, sz, 'g-')
# plt.plot(t, ez, 'c--')
# plt.plot(t, iz, 'm-')
# plt.plot(t, qz, 'y.')
# plt.plot(t, rz, 'k:')
plt.plot(t, sz + ez + iz + qz + rz, 'r-')

plt.xlabel('Time')
plt.ylabel('Total Number')
# plt.ylabel('Number of compartments')
plt.legend(['Blue', 'Red'])
# plt.legend(['Susceptible-Blue', 'Exposed-Blue', 'Infected-Blue', 'Quarantined-Blue', 'Removed-Blue'])
# plt.legend(['Susceptible-Red', 'Exposed-Red', 'Infected-Red', 'Quarantined-Red', 'Removed-Red'])
plt.show()