####################################################################
#
# The file and class for the first homework, dealing with
# radioactive decay.  The students are asked to provide a
# solution for the differential equation describing radioactive
# decay, i.e.
#
#  dN/dt = -N/tau   
#
# In the main method, at the bottom, the Euler method described in the
# lecture and implemented in DEq_Solver.py is used to solve this
# differential equation - the students will have to implement the
# relevant step there.  Here they have to provide the function
# f in the "canonical form"
#
#  dx_i/dt = f_i(x_i,t)
#
# that is being used in the lecture for describing differential equations
# of order one and their solution.  To this end, the method
#
#   dx_dt(self,x,t)
#
# needs to be modified.  
#
# In addition, in the __init__ routine the half-life given by hlife needs
# to be converted into the time constant tau. 
#
# All other input, plotting of results, etc. is already implemented 
# in the code skeleton provided.
#


import numpy


class Radioactive:
    def __init__(self,hlife):
        ### FIX THIS (Q.3c) ###
        self.tau = 0

    def dx_dt(self,x,t):
        #### FIX THIS (Q.3a) ###
        return x

    def analytical(self,x0,t):
        #### FIX THIS (Q.3b) ###
        return x0 + t

    def relative_error(self,test,t0,t1):
        res = []
        for t,x in test:
            if len(res)==0:
                x0 = x
            res.append( (t, x/self.analytical(x0,t-t0)-1) )
        return res






if __name__ == "__main__":

    from Read_Type import read_float
    from DEq_Solver import DEq_Solver

    natoms = read_float('How many atoms? ', 0)
    hlife  = read_float('Half-life of the element? ', 0.0)
    t0     = 0.
    t1     = read_float('Final time of simulation? ', 0.0)
    deltat = read_float('Size of timestep? ', 0.0, hlife)

    radioactive = Radioactive(hlife)
    deq_solver  = DEq_Solver(radioactive, 'Euler')
 
    
    x0  = numpy.array([natoms])
    result1 = deq_solver.solve(x0,t0,t1,deltat)
    ts1, xs1 = zip(*result1)

    x0 = numpy.array([natoms])
    result2 = deq_solver.solve(x0,t0,t1,deltat/10)
    ts2, xs2 = zip(*result2)

    x0 = numpy.array([natoms])
    result3 = deq_solver.solve(x0,t0,t1,5*deltat)
    ts3, xs3 = zip(*result3)




    import pylab

    # Set up first subplot
    pylab.subplot(211)
    numerical1 = pylab.semilogy(ts1,xs1,'gD',linewidth=1.0)
    numerical2 = pylab.semilogy(ts2,xs2, 'c+', linewidth=1.0)
    numerical3 = pylab.semilogy(ts3,xs3, 'bo', linewidth=1.0)

    t  = numpy.arange(t0,t1,deltat/100)
    analytical = pylab.semilogy(t, radioactive.analytical(natoms,t), 
                                linewidth=1.0, color='r')
    # Annotate the graph
    pylab.title('Simple radioactive decay')
    pylab.ylabel('number of atoms')
    pylab.setp(pylab.gca(), 'xticklabels', [])
    pylab.legend((numerical1,     numerical2,        numerical3),
           ('Numerical, dt','Numerical, dt/10','Numerical, 5*dt'))
    pylab.grid(True)


    # Set up second subplot
    pylab.subplot(212)
    pylab.ylabel('relative error (num-ana)/ana')

    errors1 = radioactive.relative_error(result1,t0,t1)
    ts1, xs1 = zip(*errors1)

    errors2 = radioactive.relative_error(result2,t0,t1)
    ts2, xs2 = zip(*errors2)

    errors3 = radioactive.relative_error(result3,t0,t1)
    ts3, xs3 = zip(*errors3)

    error1 = pylab.plot(ts1,xs1,'gD',linewidth=1.0)
    error2 = pylab.plot(ts2,xs2, 'c+', linewidth=1.0)
    error3 = pylab.plot(ts3,xs3, 'bo', linewidth=1.0)
    
    pylab.grid(True)
    pylab.xlabel('time (s)')
    # The show command displays all registered plots.
    pylab.show()