version2

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Oct 19 20:15:20 2018

@author: jenniferharber
"""

import numpy as np
import matplotlib.pyplot as plt
#from ..utils import check_random_state

#----------------------------------------------------------------------------
'>>--------------------------------Variables--------------------------------<<'

R_in   = 5    # radius of central body
M_body = 1    # mass of central body 
R_out  = 25   # radius of disc
N      = 50   # number of annuli 
H      = 1    # disc thickness
M_0    = 1    # mass entering disk 
alpha  = 2000 # viscosity constant
Factor = 0.1  # m_dot << 1
Time   = 100  # time to run simulation


#annuli = array of all the annulus (defined by calling accretion disk array so annuli = R)
# Function "accretion_disk_array" to split the disk into N annuli 
# Inputs: Number of Annuli, inner radius, outer radius
# Output: Array R of the radii of tha Annuli

#----------------------------------------------------------------------------
'>>--------------------------------Functions--------------------------------<<'

def accretion_disk_array(N, R_in, R_out):
    R = np.empty(N)
    ratio = (R_out/R_in) ** (1.0/(N-1))
    for i in range(N):
#        R[i] = (((R_out - R_in) / N) * i) + R_in   # linear progression
        R[i] = R_in * (ratio ** i)                 # geometric progression
    return R


# Function "viscous_timescale" to calculate the viscous velecity at a particular radius 
# using the formula given in (enter ref)
# Inputs: Radius
# Output: viscous velecity at that radius

def viscous_timescale(radius):
    t_viscous = 2 * np.pi * radius**(3/2)/(((H/R_out)**2) * alpha) 
    return t_viscous  


# Function "delta_visc_time" to calculate the delta viscous velecity at a given annulus 
# Inputs: Annulus
# Output: Delta Timescale

def delta_visc_time(annulus):
    delta_t = viscous_timescale(annuli[annulus+1]) - viscous_timescale(annuli[annulus])
    return int(delta_t)


# Function "m_dot" to calculate variable rate of change of mass in a particular annulus 
# Inputs: None as it cycles through each annulus
# Output: Array of lowercase m_dot changes

def m_dot():
    m_dot = np.empty((len(annuli),Time))
    for annulus in range(N):
        # For each annulus calculate the radius and viscous timescale
        # then for each time division use the timescale as the frequency of a sine wave
        
        r=annuli[annulus] #particular annulus in the annuli array 0=inner radius
        timescale = (viscous_timescale(r) / (2.0 * np.pi)) #int to try and line up waves as timescale should be a
        for t in range(Time):
            m_dot[annulus][t] = Factor * np.sin(t / timescale)
    return m_dot
  

# Function "M_dot" to calculate the overall rate of change of mass in a particular annulus 
# Inputs: None as it cycles through each annulus
# Output: Array of Capital M_dot changes
 
def M_dot():
    M_dot = np.empty((len(annuli),Time))
    #For outer annulus there is no other input than the base rate of flow into the disk (M_0)and the small change in that annuli
    for t in range(Time):
        M_dot[(len(annuli)-1)][t] = M_0 * (1.0 + m_dot[(len(annuli)-1)][t])
    #For inner annului the rate of flow is determined by the rate from the next annuli and the small change m_dot   
    #the rate from the outer annuli is time delayed by the time it takes the mass to cross that annuli  (t_offset)
    #The code cycles through the annuli from the outside (penultimate annulus) to the inside
    #As the pattern is cyclical I assume that is the time is smaller than the offset we can start the cycle again
    for i in range(len(annuli)-2,-1,-1): #backwards through array 
        t_offset_annulus = delta_visc_time(i)
        for t in range(Time):
            t_offset = t - t_offset_annulus 
            if t_offset < 0:
                t_offset = t_offset + Time
            M_dot[i][t]=M_dot[i+1][t_offset] *(1.0 + m_dot[i][t]) 
    return M_dot


#----------------------------------------------------------------------------
'>>--------------------------------Power Law--------------------------------<<'
def generate_power_law(N, dt, beta, generate_complex=False, random_state=None):
    """Generate a power-law light curve
    This uses the method from Timmer & Koenig [1]_
    Parameters
    ----------
    N : integer
        Number of equal-spaced time steps to generate
    dt : float
        Spacing between time-steps
    beta : float
        Power-law index.  The spectrum will be (1 / f)^beta
    generate_complex : boolean (optional)
        if True, generate a complex time series rather than a real time series
    random_state : None, int, or np.random.RandomState instance (optional)
        random seed or random number generator
    Returns
    -------
    x : ndarray
        the length-N
    References
    ----------
    .. [1] Timmer, J. & Koenig, M. On Generating Power Law Noise. A&A 300:707
    """
    random_state = np.check_random_state(random_state)
    dt = float(dt)
    N = int(N)

    Npos = int(N / 2)
#   Nneg = int((N - 1) / 2)
    domega = (2 * np.pi / dt / N)

    if generate_complex:
        omega = domega * np.fft.ifftshift(np.arange(N) - int(N / 2))
    else:
        omega = domega * np.arange(Npos + 1)

    x_fft = np.zeros(len(omega), dtype=complex)
    x_fft.real[1:] = random_state.normal(0, 1, len(omega) - 1)
    x_fft.imag[1:] = random_state.normal(0, 1, len(omega) - 1)

    x_fft[1:] *= (1. / omega[1:]) ** (0.5 * beta)
    x_fft[1:] *= (1. / np.sqrt(2))

    # by symmetry, the Nyquist frequency is real if x is real
    if (not generate_complex) and (N % 2 == 0):
        x_fft.imag[-1] = 0

    if generate_complex:
        x = np.fft.ifft(x_fft)
    else:
        x = np.fft.irfft(x_fft, N)

    return x


def generate_damped_RW(t_rest, tau=300., z=2.0,
                       xmean=0, SFinf=0.3, random_state=None):
    """Generate a damped random walk light curve
    This uses a damped random walk model to generate a light curve similar
    to that of a QSO [1]_.
    Parameters
    ----------
    t_rest : array_like
        rest-frame time.  Should be in increasing order
    tau : float
        relaxation time
    z : float
        redshift
    xmean : float (optional)
        mean value of random walk; default=0
    SFinf : float (optional
        Structure function at infinity; default=0.3
    random_state : None, int, or np.random.RandomState instance (optional)
        random seed or random number generator
    Returns
    -------
    x : ndarray
        the sampled values corresponding to times t_rest
    Notes
    -----
    The differential equation is (with t = time/tau):
        dX = -X(t) * dt + sigma * sqrt(tau) * e(t) * sqrt(dt) + b * tau * dt
    where e(t) is white noise with zero mean and unit variance, and
        Xmean = b * tau
        SFinf = sigma * sqrt(tau / 2)
    so
        dX(t) = -X(t) * dt + sqrt(2) * SFint * e(t) * sqrt(dt) + Xmean * dt
    References
    ----------
    .. [1] Kelly, B., Bechtold, J. & Siemiginowska, A. (2009)
           Are the Variations in Quasar Optical Flux Driven by Thermal
           Fluctuations? ApJ 698:895 (2009)
    """
    #  Xmean = b * tau
    #  SFinf = sigma * sqrt(tau / 2)
    t_rest = np.atleast_1d(t_rest)

    if t_rest.ndim != 1:
        raise ValueError('t_rest should be a 1D array')

    random_state = np.check_random_state(random_state)

    N = len(t_rest)

    t_obs = t_rest * (1. + z) / tau

    x = np.zeros(N)
    x[0] = random_state.normal(xmean, SFinf)
    E = random_state.normal(0, 1, N)

    for i in range(1, N):
        dt = t_obs[i] - t_obs[i - 1]
        x[i] = (x[i - 1]
                - dt * (x[i - 1] - xmean)
                + np.sqrt(2) * SFinf * E[i] * np.sqrt(dt))

    return x

#----------------------------------------------------------------------------
'>>--------------------------------Plotting-------------------------------<<'

#Set up the disk with N annuli
annuli= accretion_disk_array(N, R_in, R_out)

#Calculate the small periodic changes in each annuli
m_dot=m_dot() 

# code to check / show m_dot across disk - shouldn't be needed later

#y_plot=np.empty(N*Time)
#x_plot=np.empty(N*Time)
#for i in range(N*Time):
#    x_plot[i]=R_in+(i * (R_out - R_in) / (N * Time))
#    y_plot[i]=m_dot[int(i/Time)][i % Time]
#plt.plot(x_plot,y_plot)
#plt.show()

#Calculate the flow in each annuli allowing for the input to that annuli and the small changes within it
M_dot=M_dot()


# code to check / show M_dot across disk - shouldn't be needed later

y_plot=np.empty(N*Time) #annulus one for all time, annulus two for all time etc
x_plot=np.empty(N*Time)
for i in range(N*Time):
    x_plot[i]=R_in+(i * (R_out - R_in) / (N * Time)) 
    y_plot[i]=M_dot[int(i/Time)][i % Time] #
plt.plot(x_plot,y_plot)
plt.show()

# code to check / show M_dot varying in first annulus - shouldn't be needed later

y0_plot=np.empty(Time)
y1_plot=np.empty(Time)
x_plot=np.empty(Time)
print("time","annulus0","        annulus1")
for i in range(Time):
    x_plot[i]=i
    y0_plot[i]=M_dot[0][i]
    y1_plot[i]=M_dot[1][i]
    print(i,M_dot[0][i],M_dot[1][i]) 
plt.plot(x_plot,y0_plot)
plt.plot(x_plot,y1_plot)
plt.show()