# -*- coding: utf-8 -*-
# Copyright 2017 Interstellar Technologies Inc. All Rights Reserved.

from __future__ import print_function
import numpy as np
from scipy import interpolate
import matplotlib.pyplot as plt
from OpenGoddard.optimize import Problem, Guess, Condition, Dynamics

class Rocket:
    # Atmosphere Parameter
    # Use US Standard Atmosphere 1976
    stdAtmo = np.loadtxt("./11_Polar_TSTO_Taiki/US_standard_atmosphere.csv",delimiter=",",skiprows=2)
    stdAltitude = stdAtmo[:,0] * 1000.0 #converted to km -> m
    stdPressure= stdAtmo[:,2] # [Pa]
    stdDensity= stdAtmo[:,3] # [kg/m3]
    stdSoundSpeed = stdAtmo[:,4] # [m/s]
    # 線形補完用
    # 高度範囲外(<0, 86<)はfill_valueが外挿
    airPressure = interpolate.interp1d(stdAltitude, stdPressure, bounds_error = False, fill_value = (stdPressure[0], 0.0))
    airDensity = interpolate.interp1d(stdAltitude, stdDensity, bounds_error = False, fill_value = (stdDensity[0], 0.0))
    airSound = interpolate.interp1d(stdAltitude, stdSoundSpeed, bounds_error = False, fill_value = (stdSoundSpeed[0], stdSoundSpeed[-1]))

    # Drag Coefficient
    CdLog = np.loadtxt("./11_Polar_TSTO_Taiki/Cd.csv", delimiter=",", skiprows=1)
    Cd = interpolate.interp1d(CdLog[:,0], CdLog[:,1],fill_value="extrapolate")

    def __init__(self):
        # Earth Parameter
        self.GMe = 3.986004418 * 10**14  # Earth gravitational constant [m^3/s^2]
        self.Re = 6371.0 * 1000  # Earth Radius [m]
        self.g0 = 9.80665  # Gravitational acceleration on Earth surface [m/s^2]

        # Target Parameter
        self.Htarget = 561.0 * 1000  # Altitude [m]
        self.Rtarget = self.Re + self.Htarget  # Orbit Radius [m]
        self.Vtarget = np.sqrt(self.GMe / self.Rtarget)  # [m/s]

        # Launch Site Parameter
        self.lat_taiki = 42.506167 # [deg]
        self.Vt_equator = 1674.36 # [km/h]
        self.Vt_taiki = self.Vt_equator * np.cos(self.lat_taiki * np.pi / 180.0) * 1000.0 / 3600.0 # Radial Velocity of Earth Surface [m/s]
        self.inclination = 96.7 # [deg]
        self.V0 = self.Vt_taiki * np.cos(-self.inclination * np.pi / 180.0) # [m/s]
        self.H0 = 10.0  # Initial Altitude [m]

        # Structure Parameter
        # Mdryがパラメータ
        self.Mdry = [1300.0, 220.0] # Dry Mass [kg], [1st stage, 2nd stage]
        self.beta = [10.0, 15.0] # Structure Efficienty [%], [1st stage, 2nd stage]
        self.Mpayload = 100.0 # Payload Mass [kg]

        self.M0 = [self.Mdry[0] / self.beta[0] * 100.0, self.Mdry[1] / self.beta[1] * 100.0] # Initial Stage Mass [kg], [1st stage, 2nd stage]
        self.Mp = [self.M0[0] - self.Mdry[0], self.M0[1] - self.Mdry[1]] # Propellant Mass [kg], [1st stage, 2nd stage]
        self.M0[1] = self.M0[1] + self.Mpayload
        self.Minit = self.M0[0] + self.M0[1] # Initial Total Mass [kg]

        self.d = [1.8, 1.8] # Diameter [m], [1st stage, 2nd stage]
        self.A = [0.25 * self.d[0] ** 2 * np.pi, 0.25 * self.d[1] ** 2 * np.pi]  # Projected Area [m^2], [1st stage, 2nd stage]

        # Engine Parameter
        self.Cluster = 9
        self.Isp = [261.0 + 0.0, 322.0 + 0.0]  # Specific Impulse [s], [1st stage at SL, 2nd stage at vac]
        self.dth = [53.9, 53.9] # Throat Diameter [mm], [1st stage, 2nd stage]
        self.Ath = [0.25 * (self.dth[0] / 1000.0) ** 2 * np.pi, 0.25 * (self.dth[1] / 1000.0) ** 2 * np.pi] # Throat Area [m^2], [1st stage, 2nd stage]
        self.AR = [20.0, 140.0] # Area Ratio, [1st stage, 2nd stage]
        self.Ae = [self.Ath[0] * self.AR[0] * self.Cluster, self.Ath[1] * self.AR[1]] # Exit Area [m^2], [1st stage, 2nd stage]
        # =======
        self.ThrustMax = [33.3, 4.2] # Maximum Thrust [ton], [1st stage at SL, 2nd stage at vac]
        self.ThrustMax = [self.ThrustMax[0] * self.g0 * 1000.0, self.ThrustMax[1] * self.g0 * 1000.0] # [N]

        # self.ThrustLevel = 1.8 # [G] M0[0] * n G
        # self.ThrustMax = [self.M0[0] * self.ThrustLevel * self.g0, self.M0[0] * self.ThrustLevel / self.Cluster * self.g0 + self.airPressure(self.Htarget) * self.Ae[1]] # Maximum Thrust [N], [1st stage at SL, 2nd stage at vac]
        # =======

        self.refMdot = [self.ThrustMax[0] / (self.Isp[0] * self.g0), self.ThrustMax[1] / (self.Isp[1] * self.g0)] # Isp補正用参照値

        self.MaxQ = 500000.0  # Pa
        self.MaxG = 20.0  # G

def dynamics(prob, obj, section):
    R     = prob.states(0, section) # Orbit Radius [m]
    theta = prob.states(1, section) #
    Vr    = prob.states(2, section)
    Vt    = prob.states(3, section)
    m     = prob.states(4, section)
    Tr    = prob.controls(0, section)
    Tt    = prob.controls(1, section)

    g0 = obj.g0
    g = obj.g0 * (obj.Re / R)**2  # [m/s2]
    rho = obj.airDensity(R - obj.Re)
    Mach = np.sqrt(Vr**2 + Vt**2) / obj.airSound(R - obj.Re)
    Cd = obj.Cd(Mach)
    dThrust = [(obj.airPressure(obj.H0) - obj.airPressure(R - obj.Re)) * obj.Ae[0], obj.airPressure(R - obj.Re) * obj.Ae[1]]
    Isp = obj.Isp[section] + dThrust[section] / (obj.refMdot[section] * g0)

    # US standard atmosphereだと86 km以降はrho = 0でDrag = 0
    Dr = 0.5 * rho * Vr * np.sqrt(Vr**2 + Vt**2) * Cd * obj.A[section]  # [N]
    Dt = 0.5 * rho * Vt * np.sqrt(Vr**2 + Vt**2) * Cd * obj.A[section]  # [N]

    dx = Dynamics(prob, section)
    dx[0] = Vr
    dx[1] = Vt / R
    dx[2] = Tr / m - Dr / m - g + Vt**2 / R
    dx[3] = Tt / m - Dt / m - (Vr * Vt) / R
    dx[4] = - np.sqrt(Tr**2 + Tt**2) / (Isp * g0)

    return dx()

def equality(prob, obj):
    R     = prob.states_all_section(0)
    theta = prob.states_all_section(1)
    Vr    = prob.states_all_section(2)
    Vt    = prob.states_all_section(3)
    m     = prob.states_all_section(4)
    Tr    = prob.controls_all_section(0)
    Tt    = prob.controls_all_section(1)
    tf    = prob.time_final(-1)

    R0     = prob.states(0, 0)
    R1     = prob.states(0, 1)
    theta0 = prob.states(1, 0)
    theta1 = prob.states(1, 1)
    Vr0    = prob.states(2, 0)
    Vr1    = prob.states(2, 1)
    Vt0    = prob.states(3, 0)
    Vt1    = prob.states(3, 1)
    m0     = prob.states(4, 0)
    m1     = prob.states(4, 1)
    Tr0    = prob.controls(0, 0)
    Tr1    = prob.controls(0, 1)
    Tt0    = prob.controls(1, 0)
    Tt1    = prob.controls(1, 1)

    unit_R = prob.unit_states[0][0]
    unit_V = prob.unit_states[0][2]
    unit_m = prob.unit_states[0][4]

    result = Condition()

    # event condition
    result.equal(R0[0], obj.Re + obj.H0, unit=unit_R) # 初期地表
    result.equal(theta0[0], 0.0)
    result.equal(Vr0[0], 0.0, unit=unit_V)
    result.equal(Vt0[0], obj.V0 , unit=unit_V)
    result.equal(m0[0], obj.Minit, unit=unit_m) # (1st stage and 2nd stage and Payload) initial

    # knotting condition
    result.equal(m1[0], obj.M0[1], unit=unit_m) # (2nd stage + Payload) initial
    result.equal(R1[0], R0[-1], unit=unit_R)
    result.equal(theta1[0], theta0[-1])
    result.equal(Vr1[0], Vr0[-1], unit=unit_V)
    result.equal(Vt1[0], Vt0[-1], unit=unit_V)

    # Target Condition
    result.equal(R1[-1], obj.Rtarget, unit=unit_R) # Radius
    result.equal(Vr[-1], 0.0, unit=unit_V) # Radius Velocity
    result.equal(Vt[-1], obj.Vtarget, unit=unit_V)

    return result()

def inequality(prob, obj):
    R     = prob.states_all_section(0)
    theta = prob.states_all_section(1)
    Vr    = prob.states_all_section(2)
    Vt    = prob.states_all_section(3)
    m     = prob.states_all_section(4)
    Tr    = prob.controls_all_section(0)
    Tt    = prob.controls_all_section(1)
    tf    = prob.time_final(-1)

    R0     = prob.states(0, 0)
    R1     = prob.states(0, 1)
    theta0 = prob.states(1, 0)
    theta1 = prob.states(1, 1)
    Vr0    = prob.states(2, 0)
    Vr1    = prob.states(2, 1)
    Vt0    = prob.states(3, 0)
    Vt1    = prob.states(3, 1)
    m0     = prob.states(4, 0)
    m1     = prob.states(4, 1)
    Tr0    = prob.controls(0, 0)
    Tr1    = prob.controls(0, 1)
    Tt0    = prob.controls(1, 0)
    Tt1    = prob.controls(1, 1)

    rho = obj.airDensity(R - obj.Re)
    Mach = np.sqrt(Vr**2 + Vt**2) / obj.airSound(R - obj.Re)
    Cd = obj.Cd(Mach)
    Dr0 = 0.5 * rho * Vr * np.sqrt(Vr**2 + Vt**2) * Cd * obj.A[0]  # [N]
    Dt0 = 0.5 * rho * Vt * np.sqrt(Vr**2 + Vt**2) * Cd * obj.A[0]  # [N]
    Dr1 = 0.5 * rho * Vr * np.sqrt(Vr**2 + Vt**2) * Cd * obj.A[1]  # [N]
    Dt1 = 0.5 * rho * Vt * np.sqrt(Vr**2 + Vt**2) * Cd * obj.A[1]  # [N]
    g = obj.g0 * (obj.Re / R)**2  # [m/s2]

    # dynamic pressure
    q = 0.5 * rho * (Vr**2 + Vt**2)  # [Pa]
    # accelaration
    a_r0 = (Tr - Dr0) / m
    a_t0 = (Tt - Dt0) / m
    a_mag0 = np.sqrt(a_r0**2 + a_t0**2)  # [m/s2]
    a_r1 = (Tr - Dr1) / m
    a_t1 = (Tt - Dt1) / m
    a_mag1 = np.sqrt(a_r1**2 + a_t1**2)  # [m/s2]
    # Thrust
    T0 = np.sqrt(Tr0**2 + Tt0**2)
    T1 = np.sqrt(Tr1**2 + Tt1**2)
    dThrust0 = (obj.airPressure(obj.H0) - obj.airPressure(R0 - obj.Re)) * obj.Ae[0]
    dThrust1 = obj.airPressure(R1 - obj.Re) * obj.Ae[1]

    result = Condition()

    # lower bounds
    result.lower_bound(R, obj.Re, unit=prob.unit_states[0][0]) # 地表以下
    result.lower_bound(m0, obj.Mdry[0] + obj.M0[1], unit=prob.unit_states[0][4]) # 乾燥質量以下
    result.lower_bound(m1, obj.Mdry[1], unit=prob.unit_states[0][4])
    result.lower_bound(Tr, -obj.ThrustMax[1], unit=prob.unit_controls[0][0])
    result.lower_bound(Tt, -obj.ThrustMax[1], unit=prob.unit_controls[0][0])

    # upper bounds
    result.upper_bound(m0, obj.Minit, unit=prob.unit_states[0][4]) # 初期質量以上
    result.upper_bound(m1, obj.M0[1], unit=prob.unit_states[0][4])
    result.upper_bound(Tr0, obj.ThrustMax[0] + dThrust0, unit=prob.unit_controls[0][0])
    result.upper_bound(Tt0, obj.ThrustMax[0] + dThrust0, unit=prob.unit_controls[0][0])
    result.upper_bound(T0, obj.ThrustMax[0] + dThrust0, unit=prob.unit_controls[0][0])
    result.upper_bound(Tr1, obj.ThrustMax[1] + dThrust1, unit=prob.unit_controls[0][0])
    result.upper_bound(Tt1, obj.ThrustMax[1] + dThrust1, unit=prob.unit_controls[0][0])
    result.upper_bound(T1, obj.ThrustMax[1] + dThrust1, unit=prob.unit_controls[0][0])
    result.upper_bound(q, obj.MaxQ, unit = prob.unit_states[0][0])
    result.upper_bound(a_mag0, obj.MaxG * obj.g0)
    result.upper_bound(a_mag1, obj.MaxG * obj.g0)

    return result()

def cost(prob, obj):
    m1 = prob.states(4, 1)
    return -m1[-1] / prob.unit_states[1][4]

# ========================

# Program Starting Point
time_init = [0.0, 200, 800]
n = [20, 30]
num_states = [5, 5]
num_controls = [2, 2]
max_iteration = 90

flag_savefig = True
savefig_file = "./11_Polar_TSTO_Taiki/TSTO_"

# ------------------------
# set OpenGoddard class for algorithm determination
prob = Problem(time_init, n, num_states, num_controls, max_iteration)

# ------------------------
# create instance of operating object
obj = Rocket()

unit_R = obj.Re
unit_theta = 1
unit_V = np.sqrt(obj.GMe / obj.Re)
unit_m = obj.M0[0]
unit_t = unit_R / unit_V
unit_T = unit_m * unit_R / unit_t ** 2
prob.set_unit_states_all_section(0, unit_R)
prob.set_unit_states_all_section(1, unit_theta)
prob.set_unit_states_all_section(2, unit_V)
prob.set_unit_states_all_section(3, unit_V)
prob.set_unit_states_all_section(4, unit_m)
prob.set_unit_controls_all_section(0, unit_T)
prob.set_unit_controls_all_section(1, unit_T)
prob.set_unit_time(unit_t)

# ========================
# Initial parameter guess

# altitude profile
R_init = Guess.cubic(prob.time_all_section, obj.Re, 0.0, obj.Rtarget, 0.0)
# Guess.plot(prob.time_all_section, R_init, "Altitude", "time", "Altitude")
# if(flag_savefig):plt.savefig(savefig_file + "guess_alt" + ".png")
# theta
theta_init = Guess.cubic(prob.time_all_section, 0.0, 0.0, np.deg2rad(25.0), 0.0)

# velocity
Vr_init = Guess.linear(prob.time_all_section, 0.0, 0.0)
Vt_init = Guess.linear(prob.time_all_section, obj.V0, obj.Vtarget)
# Guess.plot(prob.time_all_section, V_init, "Velocity", "time", "Velocity")

# mass profile -0.6
M_init0 = Guess.cubic(prob.time_all_section, obj.Minit, 0.0, obj.Mdry[0] + obj.M0[1], 0.0)
M_init1 = Guess.cubic(prob.time_all_section, obj.M0[1], 0.0, obj.Mdry[1], 0.0)
M_init = np.hstack((M_init0, M_init1))
# Guess.plot(prob.time_all_section, M_init, "Mass", "time", "Mass")
# if(flag_savefig):plt.savefig(savefig_file + "guess_mass" + ".png")

# thrust profile
# T_init = Guess.zeros(prob.time_all_section)
Tr_init0 = Guess.cubic(prob.time[0], obj.ThrustMax[0]*9/10, 0.0, 0.0, 0.0)
Tr_init1 = Guess.cubic(prob.time[1], obj.ThrustMax[1]*9/10, 0.0, 0.0, 0.0)
Tr_init = np.hstack((Tr_init0, Tr_init1))
# Tt_init = Guess.cubic(prob.time_all_section, 0.0, 0.0, 0.0, 0.0)
Tt_init0 = Guess.cubic(prob.time[0], obj.ThrustMax[0]/10, 0.0, 0.0, 0.0)
Tt_init1 = Guess.cubic(prob.time[1], obj.ThrustMax[1]/10, 0.0, 0.0, 0.0)
Tt_init = np.hstack((Tr_init0, Tr_init1))
# Guess.plot(prob.time_all_section, T_init, "Thrust Guess", "time", "Thrust")
# if(flag_savefig):plt.savefig(savefig_file + "guess_thrust" + ".png")

# plt.show()

# ========================
# Substitution initial value to parameter vector to be optimized
# non dimensional values (Divide by scale factor)
prob.set_states_all_section(0, R_init)
prob.set_states_all_section(1, theta_init)
prob.set_states_all_section(2, Vr_init)
prob.set_states_all_section(3, Vt_init)
prob.set_states_all_section(4, M_init)
prob.set_controls_all_section(0, Tr_init)
prob.set_controls_all_section(1, Tt_init)

# ========================
# Main Process
# Assign problem to SQP solver
prob.dynamics = [dynamics, dynamics]
prob.knot_states_smooth = [False]
prob.cost = cost
# prob.cost_derivative = cost_derivative
prob.equality = equality
prob.inequality = inequality

def display_func():
    R = prob.states_all_section(0)
    theta = prob.states_all_section(1)
    m = prob.states_all_section(4)
    ts = prob.time_knots()
    tf = prob.time_final(-1)
    print("m0          : {0:.5f}".format(m[0]))
    print("mf          : {0:.5f}".format(m[-1]))
    print("mdry        : {0:.5f}".format(obj.Mdry[0]))
    print("payload     : {0:.5f}".format(m[-1] - obj.Mdry[1]))
    print("max altitude: {0:.5f}".format(R[-1] - obj.Re))
    print("MECO time   : {0:.3f}".format(ts[1]))
    print("final time  : {0:.3f}".format(tf))

prob.solve(obj, display_func, ftol=1e-8)

# ========================
# Post Process
# ------------------------
# Convert parameter vector to variable
R     = prob.states_all_section(0)
theta = prob.states_all_section(1)
Vr    = prob.states_all_section(2)
Vt    = prob.states_all_section(3)
m     = prob.states_all_section(4)
Tr    = prob.controls_all_section(0)
Tt    = prob.controls_all_section(1)
time = prob.time_update()

R0     = prob.states(0, 0)
R1     = prob.states(0, 1)
Tr0    = prob.controls(0, 0)
Tr1    = prob.controls(0, 1)
Tt0    = prob.controls(1, 0)
Tt1    = prob.controls(1, 1)

# ------------------------
# Calculate necessary variables
rho = obj.airDensity(R - obj.Re)
Mach = np.sqrt(Vr**2 + Vt**2) / obj.airSound(R - obj.Re)
Cd = obj.Cd(Mach)
Dr = 0.5 * rho * Vr * np.sqrt(Vr**2 + Vt**2) * Cd * obj.A[0]  # [N]
Dt = 0.5 * rho * Vt * np.sqrt(Vr**2 + Vt**2) * Cd * obj.A[0]  # [N]
g = obj.g0 * (obj.Re / R)**2  # [m/s2]

# dynamic pressure
q = 0.5 * rho * (Vr**2 + Vt**2)  # [Pa]
# accelaration
a_r = (Tr - Dr) / m / obj.g0
a_t = (Tt - Dt) / m / obj.g0
a_mag = np.sqrt(a_r**2 + a_t**2) / obj.g0  # [G]
# Thrust
T = np.sqrt(Tr**2 + Tt**2)

dThrust0 = (obj.airPressure(obj.H0) - obj.airPressure(R0 - obj.Re)) * obj.Ae[0]
dThrust1 = obj.airPressure(R1 - obj.Re) * obj.Ae[1]
Isp0 = obj.Isp[0] + dThrust0 / (obj.refMdot[0] * obj.g0)
Isp1 = obj.Isp[1] + dThrust1 / (obj.refMdot[1] * obj.g0)
Thrust_SL = T - np.append(dThrust0, dThrust1)
np.savetxt(savefig_file + "Thrust_Log" + ".csv", np.hstack((time, Thrust_SL, T, Tr, Tt)), delimiter=',')

# ------------------------
# Visualizetion
plt.close("all")

plt.figure()
plt.title("Altitude profile")
plt.plot(time, (R - obj.Re) / 1000, marker="o", label="Altitude")
for line in prob.time_knots():
    plt.axvline(line, color="k", alpha=0.5)
plt.grid()
plt.xlabel("time [s]")
plt.ylabel("Altitude [km]")
plt.legend(loc="best")
if(flag_savefig): plt.savefig(savefig_file + "altitude" + ".png")
np.savetxt(savefig_file + "Altitude_Log" + ".csv", np.hstack((time, (R - obj.Re))), delimiter=',')

plt.figure()
plt.title("Velocity")
plt.plot(time, Vr, marker="o", label="Vr")
plt.plot(time, Vt, marker="o", label="Vt")
for line in prob.time_knots():
    plt.axvline(line, color="k", alpha=0.5)
plt.grid()
plt.xlabel("time [s]")
plt.ylabel("Velocity [m/s]")
plt.legend(loc="best")
if(flag_savefig): plt.savefig(savefig_file + "velocity" + ".png")
np.savetxt(savefig_file + "Velocity_Log" + ".csv", np.hstack((time, Vr, Vt)), delimiter=',')

plt.figure()
plt.title("Mass")
plt.plot(time, m, marker="o", label="Mass")
for line in prob.time_knots():
    plt.axvline(line, color="k", alpha=0.5)
plt.grid()
plt.xlabel("time [s]")
plt.ylabel("Mass [kg]")
plt.legend(loc="best")
if(flag_savefig): plt.savefig(savefig_file + "mass" + ".png")
np.savetxt(savefig_file + "Mass_Log" + ".csv", np.hstack((time, m)), delimiter=',')

plt.figure()
plt.title("Acceleration")
plt.plot(time, a_r, marker="o", label="Acc r")
plt.plot(time, a_t, marker="o", label="Acc t")
plt.plot(time, a_mag, marker="o", label="Acc")
for line in prob.time_knots():
    plt.axvline(line, color="k", alpha=0.5)
plt.grid()
plt.xlabel("time [s]")
plt.ylabel("Acceleration [G]")
plt.legend(loc="best")
if(flag_savefig): plt.savefig(savefig_file + "acceleration" + ".png")

plt.figure()
plt.title("Thrust profile")
plt.plot(time, Tr / 1000, marker="o", label="Tr")
plt.plot(time, Tt / 1000, marker="o", label="Tt")
plt.plot(time, T / 1000, marker="o", label="Thrust")
plt.plot(time, Dr / 1000, marker="o", label="Dr")
plt.plot(time, Dt / 1000, marker="o", label="Dt")
plt.plot(time, m * g / 1000, marker="o", label="Gravity")
for line in prob.time_knots():
    plt.axvline(line, color="k", alpha=0.5)
plt.grid()
plt.xlabel("time [s]")
plt.ylabel("Thrust [kN]")
plt.legend(loc="best")
if(flag_savefig): plt.savefig(savefig_file + "force" + ".png")

plt.figure()
plt.title("Flight trajectory")
plt.plot(theta * obj.Re / 1000, (R - obj.Re) / 1000, marker="o", label="trajectory")
plt.grid()
plt.xlabel("Downrange [km]")
plt.ylabel("Altitude [km]")
plt.legend(loc="best")
if(flag_savefig): plt.savefig(savefig_file + "trajectory" + ".png")

plt.figure()
plt.title("DeltaThrust profile")
plt.plot(time, np.append(dThrust0, dThrust1), marker="o", label="dThrust")
for line in prob.time_knots():
    plt.axvline(line, color="k", alpha=0.5)
plt.grid()
plt.xlabel("time [s]")
plt.ylabel("dThrust [N]")
plt.legend(loc="best")
if(flag_savefig): plt.savefig(savefig_file + "dforce" + ".png")

plt.figure()
plt.title("Isp profile")
plt.plot(time, np.append(Isp0, Isp1), marker="o", label="Isp")
for line in prob.time_knots():
    plt.axvline(line, color="k", alpha=0.5)
plt.grid()
plt.xlabel("time [s]")
plt.ylabel("Isp [s]")
plt.legend(loc="best")
if(flag_savefig): plt.savefig(savefig_file + "Isp" + ".png")


#!/usr/bin/env python
'''Timing interface'''

from bisect import bisect_left, bisect_right

import numpy as np
import pandas as pd

import librosa

class TimeSlice(object):
    """
    A slice of time:  has a start time, a duration, and a reference to an Audio object.
    """

    def __init__(self, time, duration, audio, unit='s'):
        self.time = pd.to_timedelta(time, unit=unit)
        self.duration = pd.to_timedelta(duration, unit=unit)
        self.audio = audio

    def __repr__(self):
        args = self.time.delta * 1e-9, self.duration.delta * 1e-9
        return '<TimeSlice, start: {0:.2f}, duration: {1:.2f}>'.format(*args)

    def get_samples(self):
        """
        Gets the samples corresponding to this TimeSlice from the parent audio object.
        """
        start = self.time.delta * 1e-9
        duration = self.duration.delta * 1e-9
        starting_sample, ending_sample = librosa.time_to_samples(
            [start, start + duration], self.audio.sample_rate
        )

        left_offsets, right_offsets = self._get_offsets(
            starting_sample, ending_sample, self.audio.num_channels
        )

        samples = self._offset_samples(
            starting_sample,
            ending_sample,
            left_offsets,
            right_offsets,
            self.audio.num_channels,
        )

        return samples, left_offsets[0], right_offsets[0]

    def _get_offsets(self, starting_sample, ending_sample, num_channels):
        """
        Find the offset to the next zero-crossing, for each channel.
        """
        offsets = []
        for zero_index in self.audio.zero_indexes:
            index = bisect_left(zero_index, starting_sample) - 1
            if index < 0:
                starting_offset = 0
            else:
                starting_crossing = zero_index[index]
                starting_offset = starting_crossing - starting_sample

            index = bisect_left(zero_index, ending_sample)
            if index >= len(zero_index):
                ending_offset = 0
            else:
                zci = min(bisect_right(zero_index, ending_sample), len(zero_index) - 1)
                ending_crossing = zero_index[zci]
                ending_offset = ending_crossing - ending_sample

            offsets.append((starting_offset, ending_offset))

        if num_channels == 1:
            results = (offsets[0], offsets[0])
        elif num_channels == 2:
            results = (offsets[0], offsets[1])

        return results

    def _offset_samples(
        self, starting_sample, ending_sample, left_offsets, right_offsets, num_channels
    ):
        """
        Does the offset itself.
        """
        left_slice = (
            0,
            slice(starting_sample + left_offsets[0], ending_sample + left_offsets[1]),
        )
        right_slice = left_slice

        if num_channels == 2:
            right_slice = (
                1,
                slice(
                    starting_sample + right_offsets[0], ending_sample + right_offsets[1]
                ),
            )

        left_channel = self.audio.raw_samples[left_slice]
        right_channel = self.audio.raw_samples[right_slice]
        return np.array([left_channel, right_channel])

class TimingList(list):
    """
    A list of TimeSlices.
    """

    def __init__(self, name, timings, audio, unit='s'):

        super(self.__class__, self).__init__()

        self.name = name
        for (start, duration) in timings:
            time_slice = TimeSlice(start, duration, audio, unit=unit)
            self.append(time_slice)


import numpy as np
import pickle

from sklearn.isotonic import (check_increasing, isotonic_regression,
                              IsotonicRegression)

from sklearn.utils.testing import (assert_raises, assert_array_equal,
                                   assert_true, assert_false, assert_equal,
                                   assert_array_almost_equal,
                                   assert_warns_message, assert_no_warnings)
from sklearn.utils import shuffle

def test_permutation_invariance():
    # check that fit is permuation invariant.
    # regression test of missing sorting of sample-weights
    ir = IsotonicRegression()
    x = [1, 2, 3, 4, 5, 6, 7]
    y = [1, 41, 51, 1, 2, 5, 24]
    sample_weight = [1, 2, 3, 4, 5, 6, 7]
    x_s, y_s, sample_weight_s = shuffle(x, y, sample_weight, random_state=0)
    y_transformed = ir.fit_transform(x, y, sample_weight=sample_weight)
    y_transformed_s = ir.fit(x_s, y_s, sample_weight=sample_weight_s).transform(x)

    assert_array_equal(y_transformed, y_transformed_s)

def test_check_increasing_up():
    x = [0, 1, 2, 3, 4, 5]
    y = [0, 1.5, 2.77, 8.99, 8.99, 50]

    # Check that we got increasing=True and no warnings
    is_increasing = assert_no_warnings(check_increasing, x, y)
    assert_true(is_increasing)

def test_check_increasing_up_extreme():
    x = [0, 1, 2, 3, 4, 5]
    y = [0, 1, 2, 3, 4, 5]

    # Check that we got increasing=True and no warnings
    is_increasing = assert_no_warnings(check_increasing, x, y)
    assert_true(is_increasing)

def test_check_increasing_down():
    x = [0, 1, 2, 3, 4, 5]
    y = [0, -1.5, -2.77, -8.99, -8.99, -50]

    # Check that we got increasing=False and no warnings
    is_increasing = assert_no_warnings(check_increasing, x, y)
    assert_false(is_increasing)

def test_check_increasing_down_extreme():
    x = [0, 1, 2, 3, 4, 5]
    y = [0, -1, -2, -3, -4, -5]

    # Check that we got increasing=False and no warnings
    is_increasing = assert_no_warnings(check_increasing, x, y)
    assert_false(is_increasing)

def test_check_ci_warn():
    x = [0, 1, 2, 3, 4, 5]
    y = [0, -1, 2, -3, 4, -5]

    # Check that we got increasing=False and CI interval warning
    is_increasing = assert_warns_message(UserWarning, "interval",
                                         check_increasing,
                                         x, y)

    assert_false(is_increasing)

def test_isotonic_regression():
    y = np.array([3, 7, 5, 9, 8, 7, 10])
    y_ = np.array([3, 6, 6, 8, 8, 8, 10])
    assert_array_equal(y_, isotonic_regression(y))

    x = np.arange(len(y))
    ir = IsotonicRegression(y_min=0., y_max=1.)
    ir.fit(x, y)
    assert_array_equal(ir.fit(x, y).transform(x), ir.fit_transform(x, y))
    assert_array_equal(ir.transform(x), ir.predict(x))

    # check that it is immune to permutation
    perm = np.random.permutation(len(y))
    ir = IsotonicRegression(y_min=0., y_max=1.)
    assert_array_equal(ir.fit_transform(x[perm], y[perm]),
                       ir.fit_transform(x, y)[perm])
    assert_array_equal(ir.transform(x[perm]), ir.transform(x)[perm])

    # check we don't crash when all x are equal:
    ir = IsotonicRegression()
    assert_array_equal(ir.fit_transform(np.ones(len(x)), y), np.mean(y))

def test_isotonic_regression_ties_min():
    # Setup examples with ties on minimum
    x = [0, 1, 1, 2, 3, 4, 5]
    y = [0, 1, 2, 3, 4, 5, 6]
    y_true = [0, 1.5, 1.5, 3, 4, 5, 6]

    # Check that we get identical results for fit/transform and fit_transform
    ir = IsotonicRegression()
    ir.fit(x, y)
    assert_array_equal(ir.fit(x, y).transform(x), ir.fit_transform(x, y))
    assert_array_equal(y_true, ir.fit_transform(x, y))

def test_isotonic_regression_ties_max():
    # Setup examples with ties on maximum
    x = [1, 2, 3, 4, 5, 5]
    y = [1, 2, 3, 4, 5, 6]
    y_true = [1, 2, 3, 4, 5.5, 5.5]

    # Check that we get identical results for fit/transform and fit_transform
    ir = IsotonicRegression()
    ir.fit(x, y)
    assert_array_equal(ir.fit(x, y).transform(x), ir.fit_transform(x, y))
    assert_array_equal(y_true, ir.fit_transform(x, y))

def test_isotonic_regression_ties_secondary_():
    """
    Test isotonic regression fit, transform  and fit_transform
    against the "secondary" ties method and "pituitary" data from R
     "isotone" package, as detailed in: J. d. Leeuw, K. Hornik, P. Mair,
     Isotone Optimization in R: Pool-Adjacent-Violators Algorithm
    (PAVA) and Active Set Methods

    Set values based on pituitary example and
     the following R command detailed in the paper above:
    > library("isotone")
    > data("pituitary")
    > res1 <- gpava(pituitary$age, pituitary$size, ties="secondary")
    > res1$x

    `isotone` version: 1.0-2, 2014-09-07
    R version: R version 3.1.1 (2014-07-10)
    """
    x = [8, 8, 8, 10, 10, 10, 12, 12, 12, 14, 14]
    y = [21, 23.5, 23, 24, 21, 25, 21.5, 22, 19, 23.5, 25]
    y_true = [22.22222, 22.22222, 22.22222, 22.22222, 22.22222, 22.22222,
              22.22222, 22.22222, 22.22222, 24.25, 24.25]

    # Check fit, transform and fit_transform
    ir = IsotonicRegression()
    ir.fit(x, y)
    assert_array_almost_equal(ir.transform(x), y_true, 4)
    assert_array_almost_equal(ir.fit_transform(x, y), y_true, 4)

def test_isotonic_regression_reversed():
    y = np.array([10, 9, 10, 7, 6, 6.1, 5])
    y_ = IsotonicRegression(increasing=False).fit_transform(
        np.arange(len(y)), y)
    assert_array_equal(np.ones(y_[:-1].shape), ((y_[:-1] - y_[1:]) >= 0))

def test_isotonic_regression_auto_decreasing():
    # Set y and x for decreasing
    y = np.array([10, 9, 10, 7, 6, 6.1, 5])
    x = np.arange(len(y))

    # Create model and fit_transform
    ir = IsotonicRegression(increasing='auto')
    y_ = assert_no_warnings(ir.fit_transform, x, y)

    # Check that relationship decreases
    is_increasing = y_[0] < y_[-1]
    assert_false(is_increasing)

def test_isotonic_regression_auto_increasing():
    # Set y and x for decreasing
    y = np.array([5, 6.1, 6, 7, 10, 9, 10])
    x = np.arange(len(y))

    # Create model and fit_transform
    ir = IsotonicRegression(increasing='auto')
    y_ = assert_no_warnings(ir.fit_transform, x, y)

    # Check that relationship increases
    is_increasing = y_[0] < y_[-1]
    assert_true(is_increasing)

def test_assert_raises_exceptions():
    ir = IsotonicRegression()
    rng = np.random.RandomState(42)
    assert_raises(ValueError, ir.fit, [0, 1, 2], [5, 7, 3], [0.1, 0.6])
    assert_raises(ValueError, ir.fit, [0, 1, 2], [5, 7])
    assert_raises(ValueError, ir.fit, rng.randn(3, 10), [0, 1, 2])
    assert_raises(ValueError, ir.transform, rng.randn(3, 10))

def test_isotonic_sample_weight_parameter_default_value():
    # check if default value of sample_weight parameter is one
    ir = IsotonicRegression()
    # random test data
    rng = np.random.RandomState(42)
    n = 100
    x = np.arange(n)
    y = rng.randint(-50, 50, size=(n,)) + 50. * np.log(1 + np.arange(n))
    # check if value is correctly used
    weights = np.ones(n)
    y_set_value = ir.fit_transform(x, y, sample_weight=weights)
    y_default_value = ir.fit_transform(x, y)

    assert_array_equal(y_set_value, y_default_value)

def test_isotonic_min_max_boundaries():
    # check if min value is used correctly
    ir = IsotonicRegression(y_min=2, y_max=4)
    n = 6
    x = np.arange(n)
    y = np.arange(n)
    y_test = [2, 2, 2, 3, 4, 4]
    y_result = np.round(ir.fit_transform(x, y))
    assert_array_equal(y_result, y_test)

def test_isotonic_sample_weight():
    ir = IsotonicRegression()
    x = [1, 2, 3, 4, 5, 6, 7]
    y = [1, 41, 51, 1, 2, 5, 24]
    sample_weight = [1, 2, 3, 4, 5, 6, 7]
    expected_y = [1, 13.95, 13.95, 13.95, 13.95, 13.95, 24]
    received_y = ir.fit_transform(x, y, sample_weight=sample_weight)

    assert_array_equal(expected_y, received_y)

def test_isotonic_regression_oob_raise():
    # Set y and x
    y = np.array([3, 7, 5, 9, 8, 7, 10])
    x = np.arange(len(y))

    # Create model and fit
    ir = IsotonicRegression(increasing='auto', out_of_bounds="raise")
    ir.fit(x, y)

    # Check that an exception is thrown
    assert_raises(ValueError, ir.predict, [min(x) - 10, max(x) + 10])

def test_isotonic_regression_oob_clip():
    # Set y and x
    y = np.array([3, 7, 5, 9, 8, 7, 10])
    x = np.arange(len(y))

    # Create model and fit
    ir = IsotonicRegression(increasing='auto', out_of_bounds="clip")
    ir.fit(x, y)

    # Predict from  training and test x and check that min/max match.
    y1 = ir.predict([min(x) - 10, max(x) + 10])
    y2 = ir.predict(x)
    assert_equal(max(y1), max(y2))
    assert_equal(min(y1), min(y2))

def test_isotonic_regression_oob_nan():
    # Set y and x
    y = np.array([3, 7, 5, 9, 8, 7, 10])
    x = np.arange(len(y))

    # Create model and fit
    ir = IsotonicRegression(increasing='auto', out_of_bounds="nan")
    ir.fit(x, y)

    # Predict from  training and test x and check that we have two NaNs.
    y1 = ir.predict([min(x) - 10, max(x) + 10])
    assert_equal(sum(np.isnan(y1)), 2)

def test_isotonic_regression_oob_bad():
    # Set y and x
    y = np.array([3, 7, 5, 9, 8, 7, 10])
    x = np.arange(len(y))

    # Create model and fit
    ir = IsotonicRegression(increasing='auto', out_of_bounds="xyz")

    # Make sure that we throw an error for bad out_of_bounds value
    assert_raises(ValueError, ir.fit, x, y)

def test_isotonic_regression_oob_bad_after():
    # Set y and x
    y = np.array([3, 7, 5, 9, 8, 7, 10])
    x = np.arange(len(y))

    # Create model and fit
    ir = IsotonicRegression(increasing='auto', out_of_bounds="raise")

    # Make sure that we throw an error for bad out_of_bounds value in transform
    ir.fit(x, y)
    ir.out_of_bounds = "xyz"
    assert_raises(ValueError, ir.transform, x)

def test_isotonic_regression_pickle():
    y = np.array([3, 7, 5, 9, 8, 7, 10])
    x = np.arange(len(y))

    # Create model and fit
    ir = IsotonicRegression(increasing='auto', out_of_bounds="clip")
    ir.fit(x, y)

    ir_ser = pickle.dumps(ir, pickle.HIGHEST_PROTOCOL)
    ir2 = pickle.loads(ir_ser)
    np.testing.assert_array_equal(ir.predict(x), ir2.predict(x))

def test_isotonic_duplicate_min_entry():
    x = [0, 0, 1]
    y = [0, 0, 1]

    ir = IsotonicRegression(increasing=True, out_of_bounds="clip")
    ir.fit(x, y)
    all_predictions_finite = np.all(np.isfinite(ir.predict(x)))
    assert_true(all_predictions_finite)

def test_isotonic_zero_weight_loop():
    # Test from @ogrisel's issue:
    # https://github.com/scikit-learn/scikit-learn/issues/4297

    # Get deterministic RNG with seed
    rng = np.random.RandomState(42)

    # Create regression and samples
    regression = IsotonicRegression()
    n_samples = 50
    x = np.linspace(-3, 3, n_samples)
    y = x + rng.uniform(size=n_samples)

    # Get some random weights and zero out
    w = rng.uniform(size=n_samples)
    w[5:8] = 0
    regression.fit(x, y, sample_weight=w)

    # This will hang in failure case.
    regression.fit(x, y, sample_weight=w)


"""California housing dataset.

The original database is available from StatLib

    http://lib.stat.cmu.edu/datasets/

The data contains 20,640 observations on 9 variables.

This dataset contains the average house value as target variable
and the following input variables (features): average income,
housing average age, average rooms, average bedrooms, population,
average occupation, latitude, and longitude in that order.

References
----------

Pace, R. Kelley and Ronald Barry, Sparse Spatial Autoregressions,
Statistics and Probability Letters, 33 (1997) 291-297.

"""
# Authors: Peter Prettenhofer
# License: BSD 3 clause

from os.path import dirname, exists, join
from os import makedirs, remove
import tarfile

import numpy as np
import logging

import joblib

from . import get_data_home
from ._base import _convert_data_dataframe
from ._base import _fetch_remote
from ._base import _pkl_filepath
from ._base import RemoteFileMetadata
from ..utils import Bunch
from ..utils.validation import _deprecate_positional_args

# The original data can be found at:
# https://www.dcc.fc.up.pt/~ltorgo/Regression/cal_housing.tgz
ARCHIVE = RemoteFileMetadata(
    filename='cal_housing.tgz',
    url='https://ndownloader.figshare.com/files/5976036',
    checksum=('aaa5c9a6afe2225cc2aed2723682ae40'
              '3280c4a3695a2ddda4ffb5d8215ea681'))

logger = logging.getLogger(__name__)

@_deprecate_positional_args
def fetch_california_housing(*, data_home=None, download_if_missing=True,
                             return_X_y=False, as_frame=False):
    """Load the California housing dataset (regression).

    ==============   ==============
    Samples total             20640
    Dimensionality                8
    Features                   real
    Target           real 0.15 - 5.
    ==============   ==============

    Read more in the :ref:`User Guide <california_housing_dataset>`.

    Parameters
    ----------
    data_home : str, default=None
        Specify another download and cache folder for the datasets. By default
        all scikit-learn data is stored in '~/scikit_learn_data' subfolders.

    download_if_missing : bool, default=True
        If False, raise a IOError if the data is not locally available
        instead of trying to download the data from the source site.

    return_X_y : bool, default=False.
        If True, returns ``(data.data, data.target)`` instead of a Bunch
        object.

        .. versionadded:: 0.20

    as_frame : bool, default=False
        If True, the data is a pandas DataFrame including columns with
        appropriate dtypes (numeric, string or categorical). The target is
        a pandas DataFrame or Series depending on the number of target_columns.

        .. versionadded:: 0.23

    Returns
    -------
    dataset : :class:`~sklearn.utils.Bunch`
        Dictionary-like object, with the following attributes.

        data : ndarray, shape (20640, 8)
            Each row corresponding to the 8 feature values in order.
            If ``as_frame`` is True, ``data`` is a pandas object.
        target : numpy array of shape (20640,)
            Each value corresponds to the average
            house value in units of 100,000.
            If ``as_frame`` is True, ``target`` is a pandas object.
        feature_names : list of length 8
            Array of ordered feature names used in the dataset.
        DESCR : string
            Description of the California housing dataset.
        frame : pandas DataFrame
            Only present when `as_frame=True`. DataFrame with ``data`` and
            ``target``.

            .. versionadded:: 0.23

    (data, target) : tuple if ``return_X_y`` is True

        .. versionadded:: 0.20

    Notes
    -----

    This dataset consists of 20,640 samples and 9 features.
    """
    data_home = get_data_home(data_home=data_home)
    if not exists(data_home):
        makedirs(data_home)

    filepath = _pkl_filepath(data_home, 'cal_housing.pkz')
    if not exists(filepath):
        if not download_if_missing:
            raise IOError("Data not found and `download_if_missing` is False")

        logger.info('Downloading Cal. housing from {} to {}'.format(
            ARCHIVE.url, data_home))

        archive_path = _fetch_remote(ARCHIVE, dirname=data_home)

        with tarfile.open(mode="r:gz", name=archive_path) as f:
            cal_housing = np.loadtxt(
                f.extractfile('CaliforniaHousing/cal_housing.data'),
                delimiter=',')
            # Columns are not in the same order compared to the previous
            # URL resource on lib.stat.cmu.edu
            columns_index = [8, 7, 2, 3, 4, 5, 6, 1, 0]
            cal_housing = cal_housing[:, columns_index]

            joblib.dump(cal_housing, filepath, compress=6)
        remove(archive_path)

    else:
        cal_housing = joblib.load(filepath)

    feature_names = ["MedInc", "HouseAge", "AveRooms", "AveBedrms",
                     "Population", "AveOccup", "Latitude", "Longitude"]

    target, data = cal_housing[:, 0], cal_housing[:, 1:]

    # avg rooms = total rooms / households
    data[:, 2] /= data[:, 5]

    # avg bed rooms = total bed rooms / households
    data[:, 3] /= data[:, 5]

    # avg occupancy = population / households
    data[:, 5] = data[:, 4] / data[:, 5]

    # target in units of 100,000
    target = target / 100000.0

    module_path = dirname(__file__)
    with open(join(module_path, 'descr', 'california_housing.rst')) as dfile:
        descr = dfile.read()

    X = data
    y = target

    frame = None
    target_names = ["MedHouseVal", ]
    if as_frame:
        frame, X, y = _convert_data_dataframe("fetch_california_housing",
                                              data,
                                              target,
                                              feature_names,
                                              target_names)

    if return_X_y:
        return X, y

    return Bunch(data=X,
                 target=y,
                 frame=frame,
                 target_names=target_names,
                 feature_names=feature_names,
                 DESCR=descr)


"""
===================================================================
Support Vector Regression (SVR) using linear and non-linear kernels
===================================================================

Toy example of 1D regression using linear, polynomial and RBF kernels.

"""
print(__doc__)

import numpy as np
from sklearn.svm import SVR
import matplotlib.pyplot as plt

###############################################################################
# Generate sample data
X = np.sort(5 * np.random.rand(40, 1), axis=0)
y = np.sin(X).ravel()

###############################################################################
# Add noise to targets
y[::5] += 3 * (0.5 - np.random.rand(8))

###############################################################################
# Fit regression model
svr_rbf = SVR(kernel='rbf', C=1e3, gamma=0.1)
svr_lin = SVR(kernel='linear', C=1e3)
svr_poly = SVR(kernel='poly', C=1e3, degree=2)
y_rbf = svr_rbf.fit(X, y).predict(X)
y_lin = svr_lin.fit(X, y).predict(X)
y_poly = svr_poly.fit(X, y).predict(X)

###############################################################################
# look at the results
plt.scatter(X, y, c='k', label='data')
plt.hold('on')
plt.plot(X, y_rbf, c='g', label='RBF model')
plt.plot(X, y_lin, c='r', label='Linear model')
plt.plot(X, y_poly, c='b', label='Polynomial model')
plt.xlabel('data')
plt.ylabel('target')
plt.title('Support Vector Regression')
plt.legend()
plt.show()


from __future__ import print_function
import array
import pandas.msgpack as msgpack
from pandas.msgpack import ExtType

def test_pack_ext_type():
    def p(s):
        packer = msgpack.Packer()
        packer.pack_ext_type(0x42, s)
        return packer.bytes()
    assert p(b'A')        == b'\xd4\x42A'          # fixext 1
    assert p(b'AB')       == b'\xd5\x42AB'         # fixext 2
    assert p(b'ABCD')     == b'\xd6\x42ABCD'       # fixext 4
    assert p(b'ABCDEFGH') == b'\xd7\x42ABCDEFGH'   # fixext 8
    assert p(b'A'*16)     == b'\xd8\x42' + b'A'*16 # fixext 16
    assert p(b'ABC')      == b'\xc7\x03\x42ABC'        # ext 8
    assert p(b'A'*0x0123)     == b'\xc8\x01\x23\x42' + b'A'*0x0123 # ext 16
    assert p(b'A'*0x00012345) == b'\xc9\x00\x01\x23\x45\x42' + b'A'*0x00012345 # ext 32

def test_unpack_ext_type():
    def check(b, expected):
        assert msgpack.unpackb(b) == expected

    check(b'\xd4\x42A',         ExtType(0x42, b'A'))        # fixext 1
    check(b'\xd5\x42AB',        ExtType(0x42, b'AB'))       # fixext 2
    check(b'\xd6\x42ABCD',      ExtType(0x42, b'ABCD'))     # fixext 4
    check(b'\xd7\x42ABCDEFGH',  ExtType(0x42, b'ABCDEFGH')) # fixext 8
    check(b'\xd8\x42' + b'A'*16, ExtType(0x42, b'A'*16))    # fixext 16
    check(b'\xc7\x03\x42ABC',   ExtType(0x42, b'ABC'))      # ext 8
    check(b'\xc8\x01\x23\x42' + b'A'*0x0123,
          ExtType(0x42, b'A'*0x0123))                        # ext 16
    check(b'\xc9\x00\x01\x23\x45\x42' + b'A'*0x00012345,
          ExtType(0x42, b'A'*0x00012345))                   # ext 32

def test_extension_type():
    def default(obj):
        print('default called', obj)
        if isinstance(obj, array.array):
            typecode = 123 # application specific typecode
            data = obj.tostring()
            return ExtType(typecode, data)
        raise TypeError("Unknwon type object %r" % (obj,))

    def ext_hook(code, data):
        print('ext_hook called', code, data)
        assert code == 123
        obj = array.array('d')
        obj.fromstring(data)
        return obj

    obj = [42, b'hello', array.array('d', [1.1, 2.2, 3.3])]
    s = msgpack.packb(obj, default=default)
    obj2 = msgpack.unpackb(s, ext_hook=ext_hook)
    assert obj == obj2


from dnfpy.core.map2D import Map2D
import numpy as np
from sklearn.cluster import DBSCAN
import scipy.spatial.distance as dist
from dnfpyUtils.stats.statistic import Statistic

class ClusterMap(Statistic):
    """
    Params:
    "continuity" : float if different of 0.0, we assume that the cluster are continuous
        A continuous cluster allow a loss of activity during continuity seconds.
        Otherwise, the cluster is deleted
        We add the last cluster in the current coords
        Then we deduce what label labels the new cluster
        The first iteration determines the labels for the next ones
    "threshold" : threshold for activity value to be considered
    "min_sample" : how many activation are enough to be considered as a cluster
    "clustSize" : in [0,1] max size of the expected cluster
    "expectedNumberOfCluster" : as its name suggest. We will not compute anything if
    the number of activation is higher than
        (np.pi*(clustSize_/2.0)**2)*expectedNumberOfCluster
        where clustSize_ is the actual clustSize = clustSize*size

    Results:
        _data = np array (nb_clust*2) with cluster barycenter coords X,Y:
            1,2
            3,5
            3.8,3.4
    "nbOutliners" : int nb outliners found at the last compute
    if continuity > 0
    "nbNewCluster": int nb of newly build cluster

    """
    def __init__(self,name,size=0,dt=0.1,threshold=0.4,min_samples=3,
                 clustSize=0.15,sizeNpArr=1,continuity=1.0,expectedNumberOfCluster=1,
                 **kwargs):
        super(ClusterMap,self).__init__(name=name,size=size,dt=dt,threshold=threshold,
            clustSize=clustSize,min_samples=min_samples,sizeNpArr=sizeNpArr,
            continuity=continuity,expectedNumberOfCluster=expectedNumberOfCluster,
                                        **kwargs)

        self.trace = []

    def getTrace(self):
        return self.trace

    def reset(self):
        super().reset()
        self.clusters = []#cluster coords saved
        #if continuity > 0, the cluster shoul not be off for more than continuity seconds
        self.clusterOff = [] #we save the number of iteration that the cluster is off in this list
        self.setArg(nbOutliners=0)
        self.setArg(nbNewCluster=0)
        self.setArg(nbComputationEmpty=0)
        self.sumNbClusterSave = []
        self.setArg(nbClusterSum=0)
        self.setArg(maxNbAct=0)
        self.nbActList = []
        self.nbActivation = 0
        self.diffNbClusterSum = 0 #add |nbClust - expectedNumberOfCluster| at every step
        self._data = []
        self.trace = []
        self.dictOutCluster = {} #dict cluster -> output cluster
        self.outputCluster = []

    def _compute(self,size,np_arr,threshold,min_samples,clustSize_,continuity,expectedNumberOfCluster,dt,dim):
        self.toProf(size,np_arr,threshold,min_samples,clustSize_,continuity,expectedNumberOfCluster,dt,dim)

    def getMaxNbAct(self):
        if len(self.nbActList) > 0:
            return np.max(self.nbActList)
        else:
            return np.nan

    def getMeanNbAct(self):
        if len(self.nbActList) > 0:
            return np.mean(self.nbActList)
        else:
            return np.nan

    def toProf(self,size,np_arr,threshold,min_samples,clustSize_,continuity,
               expectedNumberOfCluster,dt,dim):
        maxArr = np.max(np_arr)
        coords = np.where(np_arr > maxArr/1.2)
        self.nbActivation = len(coords[0])
        #print(self.nbActivation)
        #if nbActivation > 0 and nbActivation < np_arr.shape[0]*1.6:

        #print(expectedNumberOfCluster,clustSize_)
        nbActMax = (np.pi*(clustSize_/2.0)**2)*expectedNumberOfCluster
        if self.nbActivation < nbActMax and (self.nbActivation > 0 or (continuity > 0) and (len(self.clusters)>0)):
            #print("nbActivation : %s"%self.nbActivation)
            self.nbActList.append(self.nbActivation)
            coordsArray = list(zip(coords[1],coords[0]))

            if continuity > 0:
                clusters = []
                #we add the previous valid clusters to the coordArray : hence we 'll have the same label
                for i in range(len(self._data)):
                    if not(np.any(np.isnan(self._data[i]))):
                        clust = {'id':i,'coord':self._data[i]}
                        coordsArray.insert(0,clust['coord'])
                        clusters.append(clust)
                nbClust = len(clusters)

                coordsArrayNp = np.array(coordsArray)
                distanceMat = dist.cdist(coordsArrayNp,coordsArrayNp)
                db = DBSCAN(min_samples=min_samples,eps=clustSize_).fit(distanceMat)

                #set of labels (minus outliners)
                unique_labels = set(db.labels_) - set([-1])
                #print(db.labels_)
                set_label = list(set(db.labels_))
                #print(set_label)

                clusterLabelList = []
                for i in range(len(clusters)):
                    #lab =  db.labels_[-1-i]
                    #print("index %s nb %s"%(-1-i,len(set_label)))
                    if i < len(set_label):
                        lab =  db.labels_[i]#the labels are in the same ordre as the coordArray we added
                    else:
                        print("to few label!")
                        break

                    print("i : %s. label %s , sum %s"%(i,lab,np.sum(db.labels_==lab)))
                    clusters[i]['label'] = lab
                    clusterLabelList.append( lab)

                #outliner are the labels which are -1 and not in the cluster list
                nb_outliners = len(np.where(db.labels_ == -1)[0])

                #update cluster positions
                for cluster in clusters:
                    lab = cluster['label']
                    if lab != -1 :
                        cluster['coord'] = self.__getBarycenter(coordsArrayNp,db.labels_,lab)
                    else:
                        cluster['coord'] = np.array((np.nan,)*dim)

                #check for new cluster
                nbNewCluster = 0
                for lab in set(set_label) -set(clusterLabelList):
                    barycenter = self.__getBarycenter(coordsArrayNp,db.labels_,lab)
                    #print("add cluster")
                    self._data.append(barycenter)
                    nbNewCluster += 1

                #print("nb cluster ",len(self.clusters))
                #print("nb clusterOff ",len(self.clusterOff))
                self.setArg(nbNewCluster=nbNewCluster)

                self.setArg(nbOutliners=nb_outliners)

                for cluster in clusters:
                    self._data[cluster['id']] = cluster['coord']
                print(self._data)
        elif self.nbActivation == 0:
            self.setArg(nbComputationEmpty=self.getArg("nbComputationEmpty")+1)
        else:
            #to many activation we don't compute cluster
            self._data = np.array([np.array([-1,-1])])

        self.trace.append(np.copy(self._data))

    def __getBarycenter(self,coordsArrayNp,labels,lab):
                coorBary = coordsArrayNp[np.where(labels == lab)]
                barycenter = np.mean(coorBary,axis=0)
                return barycenter

    def _onParamsUpdate(self,clustSize,sizeNpArr):
        clustSize_ = clustSize * sizeNpArr
        return dict(clustSize_=clustSize_)


#!/usr/bin/env python
"""
Create a directed graph, allowing multiple edges and self loops, from
a unix mailbox.  The nodes are email addresses with links
that point from the sender to the recievers.  The edge data
is a Python email.Message object which contains all of
the email message data.

This example shows the power of XDiGraph to hold edge data
of arbitrary Python objects (in this case a list of email messages).

By default, load the sample unix email mailbox called "unix_email.mbox".
You can load your own mailbox by naming it on the command line, eg

python unixemail.py /var/spool/mail/username

"""
__author__ = """Aric Hagberg (hagberg@lanl.gov)"""
#    Copyright (C) 2005 by
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.

import email
from email.utils import getaddresses,parseaddr
import mailbox
import sys

# unix mailbox recipe
# see http://www.python.org/doc/current/lib/module-mailbox.html
def msgfactory(fp):
    try:
        return email.message_from_file(fp)
    except email.Errors.MessageParseError:
        # Don't return None since that will stop the mailbox iterator
        return ''

if __name__ == '__main__':

    import networkx as nx
    try:
        import matplotlib.pyplot as plt
    except:
        pass

    if len(sys.argv)==1:
        filePath = "unix_email.mbox"
    else:
        filePath = sys.argv[1]

    mbox = mailbox.mbox(filePath, msgfactory) # parse unix mailbox

    G=nx.MultiDiGraph() # create empty graph

    # parse each messages and build graph
    for msg in mbox: # msg is python email.Message.Message object
        (source_name,source_addr) = parseaddr(msg['From']) # sender
        # get all recipients
        # see http://www.python.org/doc/current/lib/module-email.Utils.html
        tos = msg.get_all('to', [])
        ccs = msg.get_all('cc', [])
        resent_tos = msg.get_all('resent-to', [])
        resent_ccs = msg.get_all('resent-cc', [])
        all_recipients = getaddresses(tos + ccs + resent_tos + resent_ccs)
        # now add the edges for this mail message
        for (target_name,target_addr) in all_recipients:
            G.add_edge(source_addr,target_addr,message=msg)

    # print edges with message subject
    for (u,v,d) in G.edges_iter(data=True):
        print("From: %s To: %s Subject: %s"%(u,v,d['message']["Subject"]))

    try: # draw
        pos=nx.spring_layout(G,iterations=10)
        nx.draw(G,pos,node_size=0,alpha=0.4,edge_color='r',font_size=16)
        plt.savefig("unix_email.png")
        plt.show()
    except: # matplotlib not available
        pass


"""
Module that contains many useful utilities
for validating data or function arguments
"""
import warnings

from pandas.core.dtypes.common import is_bool

def _check_arg_length(fname, args, max_fname_arg_count, compat_args):
    """
    Checks whether 'args' has length of at most 'compat_args'. Raises
    a TypeError if that is not the case, similar to in Python when a
    function is called with too many arguments.

    """
    if max_fname_arg_count < 0:
        raise ValueError("'max_fname_arg_count' must be non-negative")

    if len(args) > len(compat_args):
        max_arg_count = len(compat_args) + max_fname_arg_count
        actual_arg_count = len(args) + max_fname_arg_count
        argument = 'argument' if max_arg_count == 1 else 'arguments'

        raise TypeError(
            "{fname}() takes at most {max_arg} {argument} "
            "({given_arg} given)".format(
                fname=fname, max_arg=max_arg_count,
                argument=argument, given_arg=actual_arg_count))

def _check_for_default_values(fname, arg_val_dict, compat_args):
    """
    Check that the keys in `arg_val_dict` are mapped to their
    default values as specified in `compat_args`.

    Note that this function is to be called only when it has been
    checked that arg_val_dict.keys() is a subset of compat_args

    """
    for key in arg_val_dict:
        # try checking equality directly with '=' operator,
        # as comparison may have been overridden for the left
        # hand object
        try:
            v1 = arg_val_dict[key]
            v2 = compat_args[key]

            # check for None-ness otherwise we could end up
            # comparing a numpy array vs None
            if (v1 is not None and v2 is None) or \
               (v1 is None and v2 is not None):
                match = False
            else:
                match = (v1 == v2)

            if not is_bool(match):
                raise ValueError("'match' is not a boolean")

        # could not compare them directly, so try comparison
        # using the 'is' operator
        except ValueError:
            match = (arg_val_dict[key] is compat_args[key])

        if not match:
            raise ValueError(("the '{arg}' parameter is not "
                              "supported in the pandas "
                              "implementation of {fname}()".
                              format(fname=fname, arg=key)))

def validate_args(fname, args, max_fname_arg_count, compat_args):
    """
    Checks whether the length of the `*args` argument passed into a function
    has at most `len(compat_args)` arguments and whether or not all of these
    elements in `args` are set to their default values.

    fname: str
        The name of the function being passed the `*args` parameter

    args: tuple
        The `*args` parameter passed into a function

    max_fname_arg_count: int
        The maximum number of arguments that the function `fname`
        can accept, excluding those in `args`. Used for displaying
        appropriate error messages. Must be non-negative.

    compat_args: OrderedDict
        A ordered dictionary of keys and their associated default values.
        In order to accommodate buggy behaviour in some versions of `numpy`,
        where a signature displayed keyword arguments but then passed those
        arguments **positionally** internally when calling downstream
        implementations, an ordered dictionary ensures that the original
        order of the keyword arguments is enforced. Note that if there is
        only one key, a generic dict can be passed in as well.

    Raises
    ------
    TypeError if `args` contains more values than there are `compat_args`
    ValueError if `args` contains values that do not correspond to those
    of the default values specified in `compat_args`

    """
    _check_arg_length(fname, args, max_fname_arg_count, compat_args)

    # We do this so that we can provide a more informative
    # error message about the parameters that we are not
    # supporting in the pandas implementation of 'fname'
    kwargs = dict(zip(compat_args, args))
    _check_for_default_values(fname, kwargs, compat_args)

def _check_for_invalid_keys(fname, kwargs, compat_args):
    """
    Checks whether 'kwargs' contains any keys that are not
    in 'compat_args' and raises a TypeError if there is one.

    """
    # set(dict) --> set of the dictionary's keys
    diff = set(kwargs) - set(compat_args)

    if diff:
        bad_arg = list(diff)[0]
        raise TypeError(("{fname}() got an unexpected "
                         "keyword argument '{arg}'".
                         format(fname=fname, arg=bad_arg)))

def validate_kwargs(fname, kwargs, compat_args):
    """
    Checks whether parameters passed to the **kwargs argument in a
    function `fname` are valid parameters as specified in `*compat_args`
    and whether or not they are set to their default values.

    Parameters
    ----------
    fname: str
        The name of the function being passed the `**kwargs` parameter

    kwargs: dict
        The `**kwargs` parameter passed into `fname`

    compat_args: dict
        A dictionary of keys that `kwargs` is allowed to have and their
        associated default values

    Raises
    ------
    TypeError if `kwargs` contains keys not in `compat_args`
    ValueError if `kwargs` contains keys in `compat_args` that do not
    map to the default values specified in `compat_args`

    """
    kwds = kwargs.copy()
    _check_for_invalid_keys(fname, kwargs, compat_args)
    _check_for_default_values(fname, kwds, compat_args)

def validate_args_and_kwargs(fname, args, kwargs,
                             max_fname_arg_count,
                             compat_args):
    """
    Checks whether parameters passed to the *args and **kwargs argument in a
    function `fname` are valid parameters as specified in `*compat_args`
    and whether or not they are set to their default values.

    Parameters
    ----------
    fname: str
        The name of the function being passed the `**kwargs` parameter

    args: tuple
        The `*args` parameter passed into a function

    kwargs: dict
        The `**kwargs` parameter passed into `fname`

    max_fname_arg_count: int
        The minimum number of arguments that the function `fname`
        requires, excluding those in `args`. Used for displaying
        appropriate error messages. Must be non-negative.

    compat_args: OrderedDict
        A ordered dictionary of keys that `kwargs` is allowed to
        have and their associated default values. Note that if there
        is only one key, a generic dict can be passed in as well.

    Raises
    ------
    TypeError if `args` contains more values than there are
    `compat_args` OR `kwargs` contains keys not in `compat_args`
    ValueError if `args` contains values not at the default value (`None`)
    `kwargs` contains keys in `compat_args` that do not map to the default
    value as specified in `compat_args`

    See Also
    --------
    validate_args : purely args validation
    validate_kwargs : purely kwargs validation

    """
    # Check that the total number of arguments passed in (i.e.
    # args and kwargs) does not exceed the length of compat_args
    _check_arg_length(fname, args + tuple(kwargs.values()),
                      max_fname_arg_count, compat_args)

    # Check there is no overlap with the positional and keyword
    # arguments, similar to what is done in actual Python functions
    args_dict = dict(zip(compat_args, args))

    for key in args_dict:
        if key in kwargs:
            raise TypeError("{fname}() got multiple values for keyword "
                            "argument '{arg}'".format(fname=fname, arg=key))

    kwargs.update(args_dict)
    validate_kwargs(fname, kwargs, compat_args)

def validate_bool_kwarg(value, arg_name):
    """ Ensures that argument passed in arg_name is of type bool. """
    if not (is_bool(value) or value is None):
        raise ValueError('For argument "{arg}" expected type bool, received '
                         'type {typ}.'.format(arg=arg_name,
                                              typ=type(value).__name__))
    return value

def validate_axis_style_args(data, args, kwargs, arg_name, method_name):
    """Argument handler for mixed index, columns / axis functions

    In an attempt to handle both `.method(index, columns)`, and
    `.method(arg, axis=.)`, we have to do some bad things to argument
    parsing. This translates all arguments to `{index=., columns=.}` style.

    Parameters
    ----------
    data : DataFrame or Panel
    arg : tuple
        All positional arguments from the user
    kwargs : dict
        All keyword arguments from the user
    arg_name, method_name : str
        Used for better error messages

    Returns
    -------
    kwargs : dict
        A dictionary of keyword arguments. Doesn't modify ``kwargs``
        inplace, so update them with the return value here.

    Examples
    --------
    >>> df._validate_axis_style_args((str.upper,), {'columns': id},
    ...                              'mapper', 'rename')
    {'columns': <function id>, 'index': <method 'upper' of 'str' objects>}

    This emits a warning
    >>> df._validate_axis_style_args((str.upper, id), {},
    ...                              'mapper', 'rename')
    {'columns': <function id>, 'index': <method 'upper' of 'str' objects>}
    """
    # TODO(PY3): Change to keyword-only args and remove all this

    out = {}
    # Goal: fill 'out' with index/columns-style arguments
    # like out = {'index': foo, 'columns': bar}

    # Start by validating for consistency
    if 'axis' in kwargs and any(x in kwargs for x in data._AXIS_NUMBERS):
        msg = "Cannot specify both 'axis' and any of 'index' or 'columns'."
        raise TypeError(msg)

    # First fill with explicit values provided by the user...
    if arg_name in kwargs:
        if args:
            msg = ("{} got multiple values for argument "
                   "'{}'".format(method_name, arg_name))
            raise TypeError(msg)

        axis = data._get_axis_name(kwargs.get('axis', 0))
        out[axis] = kwargs[arg_name]

    # More user-provided arguments, now from kwargs
    for k, v in kwargs.items():
        try:
            ax = data._get_axis_name(k)
        except ValueError:
            pass
        else:
            out[ax] = v

    # All user-provided kwargs have been handled now.
    # Now we supplement with positional arguments, emitting warnings
    # when there's ambiguity and raising when there's conflicts

    if len(args) == 0:
        pass  # It's up to the function to decide if this is valid
    elif len(args) == 1:
        axis = data._get_axis_name(kwargs.get('axis', 0))
        out[axis] = args[0]
    elif len(args) == 2:
        if 'axis' in kwargs:
            # Unambiguously wrong
            msg = ("Cannot specify both 'axis' and any of 'index' "
                   "or 'columns'")
            raise TypeError(msg)

        msg = ("Interpreting call\n\t'.{method_name}(a, b)' as "
               "\n\t'.{method_name}(index=a, columns=b)'.\nUse named "
               "arguments to remove any ambiguity. In the future, using "
               "positional arguments for 'index' or 'columns' will raise "
               " a 'TypeError'.")
        warnings.warn(msg.format(method_name=method_name,), FutureWarning,
                      stacklevel=4)
        out[data._AXIS_NAMES[0]] = args[0]
        out[data._AXIS_NAMES[1]] = args[1]
    else:
        msg = "Cannot specify all of '{}', 'index', 'columns'."
        raise TypeError(msg.format(arg_name))
    return out

def validate_fillna_kwargs(value, method, validate_scalar_dict_value=True):
    """Validate the keyword arguments to 'fillna'.

    This checks that exactly one of 'value' and 'method' is specified.
    If 'method' is specified, this validates that it's a valid method.

    Parameters
    ----------
    value, method : object
        The 'value' and 'method' keyword arguments for 'fillna'.
    validate_scalar_dict_value : bool, default True
        Whether to validate that 'value' is a scalar or dict. Specifically,
        validate that it is not a list or tuple.

    Returns
    -------
    value, method : object
    """
    from pandas.core.missing import clean_fill_method

    if value is None and method is None:
        raise ValueError("Must specify a fill 'value' or 'method'.")
    elif value is None and method is not None:
        method = clean_fill_method(method)

    elif value is not None and method is None:
        if validate_scalar_dict_value and isinstance(value, (list, tuple)):
            raise TypeError('"value" parameter must be a scalar or dict, but '
                            'you passed a "{0}"'.format(type(value).__name__))

    elif value is not None and method is not None:
        raise ValueError("Cannot specify both 'value' and 'method'.")

    return value, method


"""
:mod:`~matplotlib.gridspec` is a module which specifies the location
of the subplot in the figure.

    ``GridSpec``
        specifies the geometry of the grid that a subplot will be
        placed. The number of rows and number of columns of the grid
        need to be set. Optionally, the subplot layout parameters
        (e.g., left, right, etc.) can be tuned.

    ``SubplotSpec``
        specifies the location of the subplot in the given *GridSpec*.

"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six
from six.moves import zip

import matplotlib
rcParams = matplotlib.rcParams

import matplotlib.transforms as mtransforms

import numpy as np
import warnings

class GridSpecBase(object):
    """
    A base class of GridSpec that specifies the geometry of the grid
    that a subplot will be placed.
    """

    def __init__(self, nrows, ncols,
                 height_ratios=None, width_ratios=None):
        """
        The number of rows and number of columns of the grid need to
        be set. Optionally, the ratio of heights and widths of rows and
        columns can be specified.
        """
        #self.figure = figure
        self._nrows , self._ncols = nrows, ncols

        self.set_height_ratios(height_ratios)
        self.set_width_ratios(width_ratios)

    def get_geometry(self):
        'get the geometry of the grid, e.g., 2,3'
        return self._nrows, self._ncols

    def get_subplot_params(self, fig=None):
        pass

    def new_subplotspec(self, loc, rowspan=1, colspan=1):
        """
        create and return a SuplotSpec instance.
        """
        loc1, loc2 = loc
        subplotspec = self[loc1:loc1+rowspan, loc2:loc2+colspan]
        return subplotspec

    def set_width_ratios(self, width_ratios):
        self._col_width_ratios = width_ratios

    def get_width_ratios(self):
        return self._col_width_ratios

    def set_height_ratios(self, height_ratios):
        self._row_height_ratios = height_ratios

    def get_height_ratios(self):
        return self._row_height_ratios

    def get_grid_positions(self, fig):
        """
        return lists of bottom and top position of rows, left and
        right positions of columns.
        """
        nrows, ncols = self.get_geometry()

        subplot_params = self.get_subplot_params(fig)
        left = subplot_params.left
        right = subplot_params.right
        bottom = subplot_params.bottom
        top = subplot_params.top
        wspace = subplot_params.wspace
        hspace = subplot_params.hspace
        totWidth = right-left
        totHeight = top-bottom

        # calculate accumulated heights of columns
        cellH = totHeight/(nrows + hspace*(nrows-1))
        sepH = hspace*cellH

        if self._row_height_ratios is not None:
            netHeight = cellH * nrows
            tr = float(sum(self._row_height_ratios))
            cellHeights = [netHeight*r/tr for r in self._row_height_ratios]
        else:
            cellHeights = [cellH] * nrows

        sepHeights = [0] + ([sepH] * (nrows-1))
        cellHs = np.add.accumulate(np.ravel(list(zip(sepHeights, cellHeights))))

        # calculate accumulated widths of rows
        cellW = totWidth/(ncols + wspace*(ncols-1))
        sepW = wspace*cellW

        if self._col_width_ratios is not None:
            netWidth = cellW * ncols
            tr = float(sum(self._col_width_ratios))
            cellWidths = [netWidth*r/tr for r in self._col_width_ratios]
        else:
            cellWidths = [cellW] * ncols

        sepWidths = [0] + ([sepW] * (ncols-1))
        cellWs = np.add.accumulate(np.ravel(list(zip(sepWidths, cellWidths))))

        figTops = [top - cellHs[2*rowNum] for rowNum in range(nrows)]
        figBottoms = [top - cellHs[2*rowNum+1] for rowNum in range(nrows)]
        figLefts = [left + cellWs[2*colNum] for colNum in range(ncols)]
        figRights = [left + cellWs[2*colNum+1] for colNum in range(ncols)]

        return figBottoms, figTops, figLefts, figRights

    def __getitem__(self, key):
        """
        create and return a SuplotSpec instance.
        """
        nrows, ncols = self.get_geometry()
        total = nrows*ncols

        if isinstance(key, tuple):
            try:
                k1, k2 = key
            except ValueError:
                raise ValueError("unrecognized subplot spec")

            if isinstance(k1, slice):
                row1, row2, _ = k1.indices(nrows)
            else:
                if k1 < 0:
                    k1 += nrows
                if k1 >= nrows or k1 < 0 :
                    raise IndexError("index out of range")
                row1, row2 = k1, k1+1

            if isinstance(k2, slice):
                col1, col2, _ = k2.indices(ncols)
            else:
                if k2 < 0:
                    k2 += ncols
                if k2 >= ncols or k2 < 0 :
                    raise IndexError("index out of range")
                col1, col2 = k2, k2+1

            num1 = row1*ncols + col1
            num2 = (row2-1)*ncols + (col2-1)

        # single key
        else:
            if isinstance(key, slice):
                num1, num2, _ = key.indices(total)
                num2 -= 1
            else:
                if key < 0:
                    key += total
                if key >= total or key < 0 :
                    raise IndexError("index out of range")
                num1, num2 = key, None

        return SubplotSpec(self, num1, num2)

class GridSpec(GridSpecBase):
    """
    A class that specifies the geometry of the grid that a subplot
    will be placed. The location of grid is determined by similar way
    as the SubplotParams.
    """

    def __init__(self, nrows, ncols,
                 left=None, bottom=None, right=None, top=None,
                 wspace=None, hspace=None,
                 width_ratios=None, height_ratios=None):
        """
        The number of rows and number of columns of the
        grid need to be set. Optionally, the subplot layout parameters
        (e.g., left, right, etc.) can be tuned.
        """
        #self.figure = figure
        self.left=left
        self.bottom=bottom
        self.right=right
        self.top=top
        self.wspace=wspace
        self.hspace=hspace

        GridSpecBase.__init__(self, nrows, ncols,
                              width_ratios=width_ratios,
                              height_ratios=height_ratios)
        #self.set_width_ratios(width_ratios)
        #self.set_height_ratios(height_ratios)

    _AllowedKeys = ["left", "bottom", "right", "top", "wspace", "hspace"]

    def update(self, **kwargs):
        """
        Update the current values.  If any kwarg is None, default to
        the current value, if set, otherwise to rc.
        """

        for k, v in six.iteritems(kwargs):
            if k in self._AllowedKeys:
                setattr(self, k, v)
            else:
                raise AttributeError("%s is unknown keyword" % (k,))

        from matplotlib import _pylab_helpers
        from matplotlib.axes import SubplotBase
        for figmanager in six.itervalues(_pylab_helpers.Gcf.figs):
            for ax in figmanager.canvas.figure.axes:
                # copied from Figure.subplots_adjust
                if not isinstance(ax, SubplotBase):
                    # Check if sharing a subplots axis
                    if ax._sharex is not None and isinstance(ax._sharex, SubplotBase):
                        if ax._sharex.get_subplotspec().get_gridspec() == self:
                            ax._sharex.update_params()
                            ax.set_position(ax._sharex.figbox)
                    elif ax._sharey is not None and isinstance(ax._sharey,SubplotBase):
                        if ax._sharey.get_subplotspec().get_gridspec() == self:
                            ax._sharey.update_params()
                            ax.set_position(ax._sharey.figbox)
                else:
                    ss = ax.get_subplotspec().get_topmost_subplotspec()
                    if ss.get_gridspec() == self:
                        ax.update_params()
                        ax.set_position(ax.figbox)

    def get_subplot_params(self, fig=None):
        """
        return a dictionary of subplot layout parameters. The default
        parameters are from rcParams unless a figure attribute is set.
        """
        from matplotlib.figure import SubplotParams
        import copy
        if fig is None:
            kw = dict([(k, rcParams["figure.subplot."+k]) \
                       for k in self._AllowedKeys])
            subplotpars = SubplotParams(**kw)
        else:
            subplotpars = copy.copy(fig.subplotpars)

        update_kw = dict([(k, getattr(self, k)) for k in self._AllowedKeys])
        subplotpars.update(**update_kw)

        return subplotpars

    def locally_modified_subplot_params(self):
        return [k for k in self._AllowedKeys if getattr(self, k)]

    def tight_layout(self, fig, renderer=None, pad=1.08, h_pad=None, w_pad=None, rect=None):
        """
        Adjust subplot parameters to give specified padding.

        Parameters:

        pad : float
            padding between the figure edge and the edges of subplots, as a fraction of the font-size.
        h_pad, w_pad : float
            padding (height/width) between edges of adjacent subplots.
            Defaults to `pad_inches`.
        rect : if rect is given, it is interpreted as a rectangle
            (left, bottom, right, top) in the normalized figure
            coordinate that the whole subplots area (including
            labels) will fit into. Default is (0, 0, 1, 1).
        """

        from .tight_layout import (get_subplotspec_list,
                                   get_tight_layout_figure,
                                   get_renderer)

        subplotspec_list = get_subplotspec_list(fig.axes, grid_spec=self)
        if None in subplotspec_list:
            warnings.warn("This figure includes Axes that are not "
                          "compatible with tight_layout, so its "
                          "results might be incorrect.")

        if renderer is None:
            renderer = get_renderer(fig)

        kwargs = get_tight_layout_figure(fig, fig.axes, subplotspec_list,
                                         renderer,
                                         pad=pad, h_pad=h_pad, w_pad=w_pad,
                                         rect=rect,
                                         )

        self.update(**kwargs)

class GridSpecFromSubplotSpec(GridSpecBase):
    """
    GridSpec whose subplot layout parameters are inherited from the
    location specified by a given SubplotSpec.
    """
    def __init__(self, nrows, ncols,
                 subplot_spec,
                 wspace=None, hspace=None,
                 height_ratios=None, width_ratios=None):
        """
        The number of rows and number of columns of the grid need to
        be set. An instance of SubplotSpec is also needed to be set
        from which the layout parameters will be inherited. The wspace
        and hspace of the layout can be optionally specified or the
        default values (from the figure or rcParams) will be used.
        """
        self._wspace=wspace
        self._hspace=hspace

        self._subplot_spec = subplot_spec

        GridSpecBase.__init__(self, nrows, ncols,
                              width_ratios=width_ratios,
                              height_ratios=height_ratios)

    def get_subplot_params(self, fig=None):
        """
        return a dictionary of subplot layout parameters.
        """

        if fig is None:
            hspace = rcParams["figure.subplot.hspace"]
            wspace = rcParams["figure.subplot.wspace"]
        else:
            hspace = fig.subplotpars.hspace
            wspace = fig.subplotpars.wspace

        if self._hspace is not None:
            hspace = self._hspace

        if self._wspace is not None:
            wspace = self._wspace

        figbox = self._subplot_spec.get_position(fig, return_all=False)

        left, bottom, right, top = figbox.extents

        from matplotlib.figure import SubplotParams
        sp = SubplotParams(left=left,
                           right=right,
                           bottom=bottom,
                           top=top,
                           wspace=wspace,
                           hspace=hspace)

        return sp

    def get_topmost_subplotspec(self):
        'get the topmost SubplotSpec instance associated with the subplot'
        return self._subplot_spec.get_topmost_subplotspec()

class SubplotSpec(object):
    """
    specifies the location of the subplot in the given *GridSpec*.
    """

    def __init__(self, gridspec, num1, num2=None):
        """
        The subplot will occupy the num1-th cell of the given
        gridspec.  If num2 is provided, the subplot will span between
        num1-th cell and num2-th cell.

        The index stars from 0.
        """

        rows, cols = gridspec.get_geometry()
        total = rows*cols

        self._gridspec = gridspec
        self.num1 = num1
        self.num2 = num2

    def get_gridspec(self):
        return self._gridspec

    def get_geometry(self):
        """
        get the subplot geometry, e.g., 2,2,3. Unlike SuplorParams,
        index is 0-based
        """
        rows, cols = self.get_gridspec().get_geometry()
        return rows, cols, self.num1, self.num2

    def get_position(self, fig, return_all=False):
        """
        update the subplot position from fig.subplotpars
        """

        gridspec = self.get_gridspec()
        nrows, ncols = gridspec.get_geometry()

        figBottoms, figTops, figLefts, figRights = \
                    gridspec.get_grid_positions(fig)

        rowNum, colNum =  divmod(self.num1, ncols)
        figBottom = figBottoms[rowNum]
        figTop = figTops[rowNum]
        figLeft = figLefts[colNum]
        figRight = figRights[colNum]

        if self.num2 is not None:

            rowNum2, colNum2 =  divmod(self.num2, ncols)
            figBottom2 = figBottoms[rowNum2]
            figTop2 = figTops[rowNum2]
            figLeft2 = figLefts[colNum2]
            figRight2 = figRights[colNum2]

            figBottom = min(figBottom, figBottom2)
            figLeft = min(figLeft, figLeft2)
            figTop = max(figTop, figTop2)
            figRight = max(figRight, figRight2)

        figbox = mtransforms.Bbox.from_extents(figLeft, figBottom,
                                               figRight, figTop)

        if return_all:
            return figbox, rowNum, colNum, nrows, ncols
        else:
            return figbox

    def get_topmost_subplotspec(self):
        'get the topmost SubplotSpec instance associated with the subplot'
        gridspec = self.get_gridspec()
        if hasattr(gridspec, "get_topmost_subplotspec"):
            return gridspec.get_topmost_subplotspec()
        else:
            return self

    def __eq__(self, other):
        # check to make sure other has the attributes
        # we need to do the comparison
        if not (hasattr(other, '_gridspec') and
                hasattr(other, 'num1') and
                hasattr(other, 'num2')):
            return False
        return all((self._gridspec == other._gridspec,
                    self.num1 == other.num1,
                    self.num2 == other.num2))

    def __hash__(self):
        return (hash(self._gridspec) ^
                hash(self.num1) ^
                hash(self.num2))


#!/usr/bin/python
# Generate heatmap of points

import numpy as np
import seaborn as sns
sns.set()
import matplotlib.pyplot as plt

from heatmap_data import *

# image_name=
# im = plt.imread(image_name);
# implot = plt.imshow(im);

# Load the example flights dataset and conver to long-form
# flights_long = sns.load_dataset("flights")
# flights = flights_long.pivot("month", "year", "passengers")

def sample_kde_data(data):
    u = np.exp(data)
    z = np.sum(u)
    p = (u / z) * 1000

    xs = []
    ys = []
    for yind in xrange(len(p)):
        for xind in xrange(len(p[yind])):
            c = int(p[yind][xind])
            xs += [xind] * c
            ys += [NUM_POINTS - yind] * c

    return (np.array(xs), np.array(ys))

NUM_POINTS=25
def plot_kde(data, cmap):
    (xs, ys) = sample_kde_data(data)
    print len(xs)
    sns.kdeplot(xs, ys, cmap=cmap, shade=True, shade_lowest=False, clip=[[0,NUM_POINTS], [0, NUM_POINTS]], alpha=0.5)

# img = plt.imread("data/dqa_parts_v1/fighter-jet/fighter-jet_0000.png")
img = plt.imread("data/dqa_parts_v1/antelope/antelope_0000.png")

fig, ax = plt.subplots()
ax.imshow(img, extent=[0, NUM_POINTS, 0, NUM_POINTS])

plot_kde(neck_data3, "Blues")
# plot_kde(leg_data2, "Reds")
# plot_kde(tail_data2, "Greens")

plt.axis('off')
plt.show()

# Draw a heatmap with the numeric values in each cell
# sns.heatmap(data, cbar=False, cmap="coolwarm")
# plt.show()


"""
===========================================
Robust linear model estimation using RANSAC
===========================================

In this example we see how to robustly fit a linear model to faulty data using
the RANSAC algorithm.

"""
import numpy as np
from matplotlib import pyplot as plt

from sklearn import linear_model, datasets

n_samples = 1000
n_outliers = 50

X, y, coef = datasets.make_regression(n_samples=n_samples, n_features=1,
                                      n_informative=1, noise=10,
                                      coef=True, random_state=0)

# Add outlier data
np.random.seed(0)
X[:n_outliers] = 3 + 0.5 * np.random.normal(size=(n_outliers, 1))
y[:n_outliers] = -3 + 10 * np.random.normal(size=n_outliers)

# Fit line using all data
model = linear_model.LinearRegression()
model.fit(X, y)

# Robustly fit linear model with RANSAC algorithm
model_ransac = linear_model.RANSACRegressor(linear_model.LinearRegression())
model_ransac.fit(X, y)
inlier_mask = model_ransac.inlier_mask_
outlier_mask = np.logical_not(inlier_mask)

# Predict data of estimated models
line_X = np.arange(-5, 5)
line_y = model.predict(line_X[:, np.newaxis])
line_y_ransac = model_ransac.predict(line_X[:, np.newaxis])

# Compare estimated coefficients
print("Estimated coefficients (true, normal, RANSAC):")
print(coef, model.coef_, model_ransac.estimator_.coef_)

plt.plot(X[inlier_mask], y[inlier_mask], '.g', label='Inliers')
plt.plot(X[outlier_mask], y[outlier_mask], '.r', label='Outliers')
plt.plot(line_X, line_y, '-k', label='Linear regressor')
plt.plot(line_X, line_y_ransac, '-b', label='RANSAC regressor')
plt.legend(loc='lower right')
plt.show()