import numpy as np

# This routine both stores and provides ways of accessing a compilation of power law fits to the observed star forming sequence for different redshifts
# The compilation was constructed in the autumn of 2013 for use in Mitchell et al. (2014)

# Fits to the star forming sequence from the literature have been converted into the following functional form (used orinally in Karim et al. 2011):
# sfr / Msun / Gyr-1 = 10^(11) * (c / Gyr^-1) (Mstar / 10^(11) Msun) ^ (1+beta)
# See equation 7 in Mitchell et al. (2014) to alleviate any confusion!

# Within each dataset, there should be the following information:

# label = the source of the data
# z = the redshift of each datapoint
# beta = power law slope of the star forming sequence
# logc = log_10(c/Gyr^-1) - the normalisation of the star forming sequence at Mstar = 10^11 Msun
# c = marker colour for plotting datapoints in python
# o = marker type for plotting in python

# An example of how to extract the information in the compilation is presented at the bottom of the script

class Karim11_SF:
    """ Taken from Table 4 in Karim et al. (2011)"""
    label = "Karim et al. (2011), SF Systems"
    z = [0.3,0.5,0.7,0.9,1.1,1.4,1.8,2.25,2.75]
    beta = [-0.44, -0.42, -0.4, -0.38, -0.46, -0.3, -0.41, -0.42, -0.44]
    logc = [-1.27, -0.9, -0.67, -0.48, -0.38, -0.12, 0.1, 0.22, 0.43]
    c = "r"
    marker = "o"

class Karim11_Active:
    """ Taken from graph clicking Figure 13 of Karim et al. (2011)"""
    label = "Karim et al. (2011), Active Population"
    z = [0.3, 0.5, 0.7, 0.9, 1.1, 1.4, 1.8, 2.25, 2.75]
    beta = [-0.11, 0.03, -0.15, -0.05, -0.13, 0.01, -0.25, -0.24, -0.27]
    logc = [-0.52, -0.17, -0.2, -0.03, 0.07, 0.25, 0.3, 0.37, 0.54]
    c = "b"
    marker = "o"

class Salim07_SF:
    label = "Salim et al. (2007), SF galaxies"
    z = [0.1]
    beta = [-0.35]
    logc = [-1.18]
    c = "k"
    marker = "x"

class Peng10:
    label = "Peng et al. (2010)"
    z = [0.1]
    beta = [-0.1]
    logc = [-1.1]
    c = "b"
    marker = "x"

class Elbaz07:
    label = "Elbaz et al. (2007)"
    z = [0.1,1.0]
    beta = [-0.23,-0.1]
    logc = [-1.12,-0.267]
    c = "r"
    marker = "x"

class Daddi07:
    label = "Daddi et al. (2007)"
    z = [2.0]
    beta = [-0.1]
    logc = [0.277]
    c = "m"
    marker = "x"

class Whitaker12:
    label = "Whitaker et al. (2012)"
    z = [0.25,0.75,1.25,1.75,2.25]
    beta = [-0.33, -0.4, -0.46,-0.53,-0.59]
    logc = [-1.01,-0.57, -0.22, 0.03, 0.19]
    c = "k"
    marker = "+"

class Wang12:
    label = "Wang et al. (2012)"
    z = [0.35, 0.65, 0.9, 1.15, 1.45, 1.8]
    beta = [-0.73, -0.75, -0.72, -0.51, -0.44, -0.61]
    logc = [-1.11, -0.66, -0.44, -0.13, 0.31, 0.43]
    c = "b"
    marker = "s"

class Santini09:
    label = "Santini et al. (2009)"
    z = [0.45, 0.8, 1.25, 2.0]
    beta = [-0.3, -0.27, -0.35, -0.15]
    logc = [-0.7, -0.47, -0.42, 0.07]
    c = "r"
    marker = "D"

class Rodighiero10:
    label = "Rodighiero et al. (2010)"
    z = [0.25, 0.75, 1.25, 1.75]
    beta = [-0.24, -0.28, -0.51, -0.5]
    logc = [-0.95, -0.35, -0.05, 0.04]
    c = "k"
    marker = "*"

class Reddy12:
    label = "Reddy et al. (2012)"
    z = [2.0]
    beta = [-0.03]
    logc = [np.nan]
    marker = "+"
    c = "b"

class Sawicki12:
    label = "Sawicki et al. (2012)"
    z = [2.3]
    beta = [-0.11]
    logc = [0.074]
    c = "k"
    marker = "s"

class Oliver10:
    label = "Oliver et al. (2010)"
    z = [0.1, 0.25, 0.35, 0.45, 0.55, 0.7, 0.9, 1.125, 1.375, 1.75]
    beta = [0.0, 0.01, -0.16, -0.12, -0.2, -0.26, -0.3, -0.32, -0.3, 0.17]
    logc = [-1.08, -0.93, -0.96, -0.84, -0.75, -0.61, -0.37, -0.17, 0.09, 0.04]
    c = "b"
    marker = "D"

class Magdis10:
    label = "Magdis et al. (2010)"
    z = [3.0]
    beta = [-0.09]
    logc = [0.54]
    c = "r"
    marker = "*"

class Lin12:
    label = "Lin et al. (2012)"
    z = [2.0]
    beta = [-0.26]
    logc = [np.nan]
    marker = "^"
    c = "k"

class Labbe10:
    label = "Labbe et al. (2010)"
    z = [7.0]
    beta = [-0.06]
    logc = [0.152]
    marker = "s"
    c = "g"

class Koyama13:
    label = "Koyama et al. (2013)"
    z = [0.4,0.8, 2.2]
    beta = [-0.36,-0.53,-0.51]
    logc = [-1.25,-0.82, -0.252]
    marker = "D"
    c = "g"

class Hathi13:
    label = "Hathi et al. (2013)"
    z = [1.6,2.2,2.6,4.5]
    beta = [-0.1,-0.1,-0.1,-0.1]
    logc = [0.982, 0.64, -0.9, 0.87]
    marker = "p"
    c = "r"

class Greene13:
    label = "Greene et al. (2013)"
    z = [0.75]
    beta = [-0.2]
    logc = [-0.85]
    marker = "p"
    c = "m"

class Drory08:
    label = "Drory et al. (2008)"
    z = [0.5,1.0,1.5,2.0,2.75,3.5]#,4.5]
    beta = [-0.31, -0.44, -0.48, -0.39, -0.41, -0.37]#, -0.34]
    logc = [-0.8, -0.82, -0.66, -0.4, -0.37, -0.05]#, -0.61]
    marker = "1"
    c = "k"

class Elbaz11:
    label = "Elbaz et al. (2011)"
    z =    [0.1, 0.146, 0.217, 0.452, 0.666, 0.748, 1.044, 1.125, 1.574, 1.646, 1.952, 2.238, 2.462, 2.023, 2.891]
    beta = [0.0, 0.0,   0.0,   0.0,   0.0,   0.0,   0.0,   0.0,   0.0,   0.0,   0.0,   0.0,   0.0,   0.0,   0.0]
    logc = [-0.6,-0.93, -0.78, -0.76, -0.49, -0.59, -0.52, -0.25, 0.28,  0.2,   0.23,  0.28,  0.44,  0.78, 0.49]
    marker = ">"
    c = "k"

class Bouwens12:
    label = "Bouwens et al. (2012)"
    z = [4.0]
    beta = [-0.27]
    logc = [0.51]
    marker = "+"
    c = "r"

# Very stupid code but it should work!

def Get_Sequence(name):
    if name == "Karim11_SF":
        z = Karim11_SF.z; beta = Karim11_SF.beta; logc = Karim11_SF.logc; label = Karim11_SF.label; c = Karim11_SF.c; marker = Karim11_SF.marker
    elif name == "Karim11_Active":
        z = Karim11_Active.z; beta = Karim11_Active.beta; logc = Karim11_Active.logc; label = Karim11_Active.label; c = Karim11_Active.c; marker = Karim11_Active.marker
    elif name == "Salim07_SF":
        z = Salim07_SF.z; beta = Salim07_SF.beta; logc = Salim07_SF.logc; label = Salim07_SF.label; c = Salim07_SF.c; marker = Salim07_SF.marker
    elif name == "Peng10":
        z = Peng10.z; beta = Peng10.beta; logc = Peng10.logc; label = Peng10.label; c = Peng10.c; marker = Peng10.marker
    elif name == "Elbaz07":
        z = Elbaz07.z; beta = Elbaz07.beta; logc = Elbaz07.logc; label = Elbaz07.label; c = Elbaz07.c; marker = Elbaz07.marker
    elif name == "Daddi07":
        z = Daddi07.z; beta = Daddi07.beta; logc = Daddi07.logc; label = Daddi07.label; c = Daddi07.c; marker = Daddi07.marker
    elif name == "Whitaker12":
        z = Whitaker12.z; beta = Whitaker12.beta; logc = Whitaker12.logc; label = Whitaker12.label; c = Whitaker12.c; marker = Whitaker12.marker
    elif name == "Wang12":
        z = Wang12.z; beta = Wang12.beta; logc = Wang12.logc; label = Wang12.label; c = Wang12.c; marker = Wang12.marker
    elif name == "Santini09":
        z = Santini09.z; beta = Santini09.beta; logc = Santini09.logc; label = Santini09.label; c = Santini09.c; marker = Santini09.marker
    elif name == "Rodighiero10":
        z = Rodighiero10.z; beta = Rodighiero10.beta; logc = Rodighiero10.logc; label = Rodighiero10.label; c = Rodighiero10.c; marker = Rodighiero10.marker
    elif name == "Reddy12":
        z = Reddy12.z; beta = Reddy12.beta; logc = Reddy12.logc; label = Reddy12.label; c = Reddy12.c; marker = Reddy12.marker
    elif name == "Sawicki12":
        z = Sawicki12.z; beta = Sawicki12.beta; logc = Sawicki12.logc; label = Sawicki12.label; c = Sawicki12.c; marker = Sawicki12.marker
    elif name == "Magdis12":
        z = Magdis12.z; beta = Magdis12.beta; logc = Magdis12.logc; label = Magdis12.label; c = Magdis12.c; marker = Magdis12.marker
    elif name == "Oliver10":
        z = Oliver10.z; beta = Oliver10.beta; logc = Oliver10.logc; label = Oliver10.label; c = Oliver10.c; marker = Oliver10.marker
    elif name == "Lin12":
        z = Lin12.z; beta = Lin12.beta; logc = Lin12.logc; label = Lin12.label; c = Lin12.c; marker = Lin12.marker
    elif name == "Labbe10":
        z = Labbe10.z; beta = Labbe10.beta; logc = Labbe10.logc; label = Labbe10.label; c = Labbe10.c; marker = Labbe10.marker
    elif name == "Koyama13":
        z = Koyama13.z; beta = Koyama13.beta; logc = Koyama13.logc; label = Koyama13.label; c = Koyama13.c; marker = Koyama13.marker
    elif name == "Huang12":
        z = Huang12.z; beta = Huang12.beta; logc = Huang12.logc; label = Huang12.label; c = Huang12.c; marker = Huang12.marker
    elif name == "Huang12_blue":
        z = Huang12_blue.z; beta = Huang12_blue.beta; logc = Huang12_blue.logc; label = Huang12_blue.label; c = Huang12_blue.c; marker = Huang12_blue.marker
    elif name == "Hathi13":
        z = Hathi13.z; beta = Hathi13.beta; logc = Hathi13.logc; label = Hathi13.label; c = Hathi13.c; marker = Hathi13.marker
    elif name == "Greene13":
        z = Greene13.z; beta = Greene13.beta; logc = Greene13.logc; label = Greene13.label; c = Greene13.c; marker = Greene13.marker
    elif name == "Drory08":
        z = Drory08.z; beta = Drory08.beta; logc = Drory08.logc; label = Drory08.label; c = Drory08.c; marker = Drory08.marker
    elif name == "Elbaz11":
        z = Elbaz11.z; beta = Elbaz11.beta; logc = Elbaz11.logc; label = Elbaz11.label; c = Elbaz11.c; marker = Elbaz11.marker
    elif name == "Bouwens12":
        z = Bouwens12.z; beta = Bouwens12.beta; logc = Bouwens12.logc; label = Bouwens12.label; c = Bouwens12.c; marker = Bouwens12.marker
    else:
        print "Error, name = ", name," is not recognised"
        quit()
    return z, beta, logc, label, c, marker

list = ["Daddi07","Elbaz07","Salim07_SF","Santini09","Labbe10","Oliver10","Peng10","Rodighiero10","Elbaz11","Karim11_SF","Karim11_Active","Bouwens12","Koyama13","Lin12","Reddy12","Sawicki12","Wang12","Whitaker12"]

# Sources that were chosen not to be presented within Mitchell et al. (2014)
excluded = ["Hathi13","Greene13","Drory08"]

# Example usage of the script

if __name__=='__main__':
    from pylab import *; from scipy import integrate

    # example cosmological parameters (used to calculate lookback time)
    h = 0.7; omm = 0.25; oml = 0.75

    H = lambda z: 100 * h * (omm*(1.0+z)**3.0 + oml)**0.5
    age_int = lambda z: 1.0/((1+z)*H(z))

    # calculate the age of the universe at z=0
    age_universe = integrate.quad(age_int,float(0.),np.inf)[0] * 979.16

    # Desired list of sources of observational data
    name_list = ["Daddi07","Elbaz07","Salim07_SF","Santini09","Labbe10","Oliver10","Peng10","Rodighiero10","Elbaz11","Karim11_SF","Karim11_Active","Bouwens12","Koyama13","Lin12","Reddy12","Sawicki12","Wang12","Whitaker12"]

    # Get the data and plot the estimated slope and normalisation as a function of lookback time

    for name in name_list:
        z,beta,logc,label = Get_Sequence(name)
        
        # calculate the lookback times corresponding to each redshift
        tage = np.zeros_like(z)
        for n in range(len(z)):
            tage[n] = integrate.quad(age_int,float(z[n]),np.inf)[0] * 979.16 #Gyr
            tlb = age_universe - tage
        
        # plot the star forming sequence power-law slope as a function of lookback time
        subplot(211)
        scatter(tlb, beta,label=label)
        # plot the star forming sequence normalisation as a function of lookback time
        subplot(212)
        scatter(tlb,logc)
    
    xlim((0.0,13.5)); ylim((-2.0,2.0))
    legend()
    show()