import numpy as np
import openmdao.api as om

class SIRvec(om.ExplicitComponent):
    """Basic epidemiological infection model
       S (suceptible), I (infected), R (recovered).
    """

    def initialize(self):
        self.options.declare('num_nodes', types=int)
        self.options.declare('num_runs', types=int)

    def setup(self):

        nn = self.options['num_nodes']
        nr = self.options['num_runs']

        # States
        self.add_input('S',
               val=np.zeros((nr, nn)))

        self.add_input('I',
               val=np.zeros((nr, nn)))

        self.add_input('R',
               val=np.zeros((nr, nn)))

        # ROCs
        self.add_output('Sdot', val=np.zeros((nr, nn)))
        self.add_output('Idot', val=np.zeros((nr, nn)))
        self.add_output('Rdot', val=np.zeros((nr, nn)))

        # Params

        self.add_input('beta',
               val = np.zeros((nr, nn)))

        self.add_input('gamma',
               val = np.zeros((nr, nn)))

        #self.add_output('max_I', val=0.0)

        #arange = np.arange(self.options['num_nodes'], dtype=int)
        arange = np.arange(nn + nr, dtype=int)

        self.declare_partials('Sdot', ['beta', 'gamma', 'S', 'I', 'R'], rows=arange, cols=arange)

        self.declare_partials('Idot', ['beta', 'gamma', 'S', 'I', 'R'], rows=arange, cols=arange)

        self.declare_partials('Rdot', ['gamma', 'I'], rows=arange, cols=arange)

        #self.declare_partials('max_I', 'I')

    def compute(self, inputs, outputs):

        S = inputs['S']
        I = inputs['I']
        R = inputs['R']
        gamma = inputs['gamma']
        beta = inputs['beta']

        # time derivatives of the states of an SIR model
        # substitution dynamic control 'beta' for constant 'beta'.
        outputs['Sdot'] = - beta * S * I
        outputs['Idot'] = beta * S * I - gamma * I
        outputs['Rdot'] = gamma * I


    def compute_partials(self, inputs, jacobian):

        S = inputs['S']
        I = inputs['I']
        R = inputs['R']
        gamma = inputs['gamma']
        beta = inputs['beta']

        # derivatives of the ODE equations of state
        jacobian['Sdot', 'S'] = (-I * beta).flatten()
        jacobian['Sdot', 'I'] = (-S * beta).flatten()
        jacobian['Sdot', 'beta'] = ( - S * I).flatten()


        jacobian['Idot', 'gamma'] = (-I).flatten()
        jacobian['Idot', 'S'] = I * (beta).flatten()
        jacobian['Idot', 'I'] = S * (beta - gamma).flatten()
        jacobian['Idot', 'beta'] = (S * I).flatten()

        jacobian['Rdot', 'gamma'] = (I).flatten()
        jacobian['Rdot', 'I'] = (gamma).flatten()


if __name__ == '__main__':
    np.random.seed(0)

    p = om.Problem()
    p.model = om.Group()
    nr = 4
    n = 35
    p.model.add_subsystem('test', SIRvec(num_nodes=n, num_runs=nr), promotes=['*'])
    p.setup()
    np.random.seed(0)
    p['S'] = np.random.uniform(1, 1000, (nr, n))
    p['I'] = np.random.uniform(1, 1000, (nr, n))
    p['R'] = np.random.uniform(1, 1000, (nr, n))

    p['beta'] = np.random.uniform(0, 2, (nr, n))
    p['gamma'] = np.random.uniform(0, 2, (nr, n))

    #p['t'] = np.linspace(0, 100, n)
    
    p.run_model()

    x = p.check_partials(compact_print=True)