import numpy as np
import warnings
import sys
try:
    import pycuda.autoinit
    import pycuda.driver as drv
    from pycuda import gpuarray
    from pycuda.compiler import SourceModule
    import pycuda.cumath
    pycuda_available = True
except:
    pycuda_available = False

use_gpu = False

default_blocksize = 128

def enable_gpu(enable=True):
    """Sets the use_gpu flag to enable/disable the use of CUDA kernels.

    Args:
        enable (bool): Set use_gpu flag to this value (default=True).
    """
    # todo: do gpuarrays require cuda toolkit? otherwise distinguish if only gpuarrays work but no kernel compilation
    global use_gpu

    if use_gpu == enable:
        return

    if not enable:
        use_gpu = False
        sys.stdout.write('Disabling GPU usage.\n')
        sys.stdout.flush()
        return

    if not pycuda_available:
        warnings.warn("Unable to import PyCuda - fall back to CPU mode")
        use_gpu = False
        sys.stdout.write('Disabling GPU usage.\n')
        sys.stdout.flush()
    else:
        try:
            test_fun = SourceModule("""
                __global__ void test_kernel(float *x)
                {
                    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
                    x[i] = 2 * i;
                }""").get_function("test_kernel")
            x = gpuarray.to_gpu(np.zeros(128, dtype=np.float32))
            test_fun(x, block=(128,1,1), grid=(1,1))
            assert np.allclose(x.get(), 2 * np.arange(128))
            use_gpu = True
            sys.stdout.write('Enabling GPU usage.\n')
            sys.stdout.flush()
            return

        except Exception as e:
            warnings.warn("CUDA test kernel failed - fall back to CPU mode")
            print(e)
            use_gpu = False
            sys.stdout.write('Disabling GPU usage.\n')
            sys.stdout.flush()

# The following cuda kernel multiplies the coupling matrix to a vector. It is based on linear interpolation of
# the lookup table.
#
# input arguments of the cuda kernel:
# n (np.uint32):  n[i] contains the mutlipole multi-index with regard to self.l_max and self.m_max of the i-th
#                 entry of a system vector
# m (np.float32): m[i] contains the multipole order with regard to particle.l_max and particle.m_max
# x_pos (np.float32): x_pos[i] contains the respective particle x-position
# y_pos (np.float32): y_pos[i] contains the respective particle y-position
# z_pos (np.float32): z_pos[i] contains the respective particle z-position
# re_lookup_pl (np.float32): the real part of the lookup table for the z1+z2 part of the Sommerfeld integral,
#                            in the format (rho, sum_z, n1, n2)
# im_lookup_pl (np.float32): the imaginary part of the lookup table for the z1+z2 part of the Sommerfeld
#                            integral, in the format (rho, sum_z, n1, n2)
# re_lookup_mn (np.float32): the real part of the lookup table for the z1-z2 part of the Sommerfeld integral,
#                            in the format (rho, diff_z, n1, n2)
# im_lookup_mn (np.float32): the imaginary part of the lookup table for the z1-z2 part of the Sommerfeld
#                            integral, in the format (rho, diff_z, n1, n2)
# re_in_vec (np.float32): the real part of the vector to be multiplied with the coupling matrix
# im_in_vec (np.float32): the imaginary part of the vector to be multiplied with the coupling matrix
# re_result_vec (np.float32): the real part of the vector into which the result is written
# im_result_vec (np.float32): the imaginary part of the vector into which the result is written
linear_volume_lookup_source = """
    #define BLOCKSIZE %i
    #define NUMBER_OF_UNKNOWNS %i
    #define Z_ARRAY_LENGTH %i
    #define MIN_RHO %f
    #define MIN_Z_SUM %f
    #define MIN_Z_DIFF %f
    #define LOOKUP_RESOLUTION %f
    __global__ void coupling_kernel(const int *n, const float *m, const float *x_pos, const float *y_pos,
                                    const float *z_pos, const float *re_lookup_pl, const float *im_lookup_pl,
                                    const float *re_lookup_mn, const float *im_lookup_mn,
                                    const float *re_in_vec, const float *im_in_vec,
                                    float  *re_result, float  *im_result)
    {
        unsigned int i1 = blockIdx.x * blockDim.x + threadIdx.x;
        if(i1 >= NUMBER_OF_UNKNOWNS) return;

        const float x1 = x_pos[i1];
        const float y1 = y_pos[i1];
        const float z1 = z_pos[i1];

        const int n1 = n[i1];
        const float m1 = m[i1];

        re_result[i1] = 0.0;
        im_result[i1] = 0.0;

        for (int i2=0; i2<NUMBER_OF_UNKNOWNS; i2++)
        {
            float x21 = x1 - x_pos[i2];
            float y21 = y1 - y_pos[i2];
            float sz21 = z1 + z_pos[i2];
            float dz21 = z1 - z_pos[i2];

            const int n2 = n[i2];
            const float m2 = m[i2];

            float rho = sqrt(x21*x21+y21*y21);
            float phi = atan2(y21,x21);

            int rho_idx = (int) floor((rho - MIN_RHO) / LOOKUP_RESOLUTION);
            float rho_w = (rho - MIN_RHO) / LOOKUP_RESOLUTION - floor((rho - MIN_RHO) / LOOKUP_RESOLUTION);

            int sz_idx = (int) floor((sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION);
            float sz_w = (sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION - floor((sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION);

            int dz_idx = (int) floor((dz21 - MIN_Z_DIFF) / LOOKUP_RESOLUTION);
            float dz_w = (dz21 - MIN_Z_DIFF) / LOOKUP_RESOLUTION
                            - floor((dz21 - MIN_Z_DIFF) /  LOOKUP_RESOLUTION);

            int idx_rho_sz = rho_idx * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                +  sz_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rho_szpl1 = rho_idx * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                + (sz_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rhopl1_sz = (rho_idx + 1) * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                + sz_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rhopl1_szpl1 = (rho_idx + 1) * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                    + (sz_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rho_dz = rho_idx * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE + dz_idx * BLOCKSIZE * BLOCKSIZE
                                + n1 * BLOCKSIZE + n2;

            int idx_rho_dzpl1 = rho_idx * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                + (dz_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rhopl1_dz = (rho_idx + 1) * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                + dz_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rhopl1_dzpl1 = (rho_idx + 1) * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                    + (dz_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            float f_rhoi = re_lookup_pl[idx_rho_sz] * (1 - sz_w) + re_lookup_pl[idx_rho_szpl1] * sz_w;
            float f_rhoipl1 = re_lookup_pl[idx_rhopl1_sz] * (1 - sz_w) + re_lookup_pl[idx_rhopl1_szpl1] * sz_w;
            float re_si_pl = f_rhoi * (1 - rho_w) + f_rhoipl1 * rho_w;

            f_rhoi = im_lookup_pl[idx_rho_sz] * (1 - sz_w) + im_lookup_pl[idx_rho_szpl1] * sz_w;
            f_rhoipl1 = im_lookup_pl[idx_rhopl1_sz] * (1 - sz_w) + im_lookup_pl[idx_rhopl1_szpl1] * sz_w;
            float im_si_pl = f_rhoi * (1 - rho_w) + f_rhoipl1 * rho_w;

            f_rhoi = re_lookup_mn[idx_rho_dz] * (1 - dz_w) + re_lookup_mn[idx_rho_dzpl1] * dz_w;
            f_rhoipl1 = re_lookup_mn[idx_rhopl1_dz] * (1 - dz_w) + re_lookup_mn[idx_rhopl1_dzpl1] * dz_w;
            float re_si_mn = f_rhoi * (1 - rho_w) + f_rhoipl1 * rho_w;

            f_rhoi = im_lookup_mn[idx_rho_dz] * (1 - dz_w) + im_lookup_mn[idx_rho_dzpl1] * dz_w;
            f_rhoipl1 = im_lookup_mn[idx_rhopl1_dz] * (1 - dz_w) + im_lookup_mn[idx_rhopl1_dzpl1] * dz_w;
            float im_si_mn = f_rhoi * (1 - rho_w) + f_rhoipl1 * rho_w;

            float re_eimphi = cosf((m2 - m1) * phi);
            float im_eimphi = sinf((m2 - m1) * phi);

            float re_w = re_eimphi * (re_si_pl + re_si_mn) - im_eimphi * (im_si_pl + im_si_mn);
            float im_w = im_eimphi * (re_si_pl + re_si_mn) + re_eimphi * (im_si_pl + im_si_mn);

            re_result[i1] += re_w * re_in_vec[i2] - im_w * im_in_vec[i2];
            im_result[i1] += re_w * im_in_vec[i2] + im_w * re_in_vec[i2];
        }
    }"""

# The following cuda kernel multiplies the coupling matrix to a vector. It is based on cubic Hermite spline interpolation
# of the lookup table.
#
# input arguments of the cuda kernel:
# n (np.uint32):  n[i] contains the mutlipole multi-index with regard to self.l_max and self.m_max of the i-th
#                 entry of a system vector
# m (np.float32): m[i] contains the multipole order with regard to particle.l_max and particle.m_max
# x_pos (np.float32): x_pos[i] contains the respective particle x-position
# y_pos (np.float32): y_pos[i] contains the respective particle y-position
# z_pos (np.float32): z_pos[i] contains the respective particle z-position
# re_lookup_pl (np.float32): the real part of the lookup table for the z1+z2 part of the Sommerfeld integral,
#                            in the format (rho, sum_z, n1, n2)
# im_lookup_pl (np.float32): the imaginary part of the lookup table for the z1+z2 part of the Sommerfeld
#                            integral, in the format (rho, sum_z, n1, n2)
# re_lookup_mn (np.float32): the real part of the lookup table for the z1-z2 part of the Sommerfeld integral,
#                            in the format (rho, diff_z, n1, n2)
# im_lookup_mn (np.float32): the imaginary part of the lookup table for the z1-z2 part of the Sommerfeld
#                            integral, in the format (rho, diff_z, n1, n2)
# re_in_vec (np.float32): the real part of the vector to be multiplied with the coupling matrix
# im_in_vec (np.float32): the imaginary part of the vector to be multiplied with the coupling matrix
# re_result_vec (np.float32): the real part of the vector into which the result is written
# im_result_vec (np.float32): the imaginary part of the vector into which the result is written
cubic_volume_lookup_source = """
    #define BLOCKSIZE %i
    #define NUMBER_OF_UNKNOWNS %i
    #define Z_ARRAY_LENGTH %i
    #define MIN_RHO %f
    #define MIN_Z_SUM %f
    #define MIN_Z_DIFF %f
    #define LOOKUP_RESOLUTION %f

    __device__ int lookup_index(int const i_r, int const i_z, int const n1, int const n2)
    {
        return i_r * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE + i_z * BLOCKSIZE * BLOCKSIZE + n1 *  BLOCKSIZE + n2;
    }

    __device__ float cubic_interpolation(float const w, float const lookup_imn1, float const lookup_i,
                                            float const lookup_ipl1, float const lookup_ipl2)
    {
        return ((-w*w*w+2*w*w-w) * lookup_imn1 + (3*w*w*w-5*w*w+2) * lookup_i + (-3*w*w*w+4*w*w+w) * lookup_ipl1
                + (w*w*w-w*w) * lookup_ipl2) / 2;
    }

    __device__ float lookup_2D(int const i_rho, float const w_rho, int const i_z, float const w_z, int const n1,
                                int const n2, float const *lookup)
    {
        int i_zmn1 = lookup_index(i_rho-1, i_z-1, n1, n2);
        int i = lookup_index(i_rho-1, i_z, n1, n2);
        int i_zpl1 = lookup_index(i_rho-1, i_z+1, n1, n2);
        int i_zpl2 = lookup_index(i_rho-1, i_z+2, n1, n2);
        float lookup_rmn1 = cubic_interpolation(w_z, lookup[i_zmn1], lookup[i], lookup[i_zpl1], lookup[i_zpl2]);

        i_zmn1 = lookup_index(i_rho, i_z-1, n1, n2);
        i = lookup_index(i_rho, i_z, n1, n2);
        i_zpl1 = lookup_index(i_rho, i_z+1, n1, n2);
        i_zpl2 = lookup_index(i_rho, i_z+2, n1, n2);
        float lookup_r = cubic_interpolation(w_z, lookup[i_zmn1], lookup[i], lookup[i_zpl1], lookup[i_zpl2]);

        i_zmn1 = lookup_index(i_rho+1, i_z-1, n1, n2);
        i = lookup_index(i_rho+1, i_z, n1, n2);
        i_zpl1 = lookup_index(i_rho+1, i_z+1, n1, n2);
        i_zpl2 = lookup_index(i_rho+1, i_z+2, n1, n2);
        float lookup_rpl1 = cubic_interpolation(w_z, lookup[i_zmn1], lookup[i], lookup[i_zpl1], lookup[i_zpl2]);

        i_zmn1 = lookup_index(i_rho+2, i_z-1, n1, n2);
        i = lookup_index(i_rho+2, i_z, n1, n2);
        i_zpl1 = lookup_index(i_rho+2, i_z+1, n1, n2);
        i_zpl2 = lookup_index(i_rho+2, i_z+2, n1, n2);
        float lookup_rpl2 = cubic_interpolation(w_z, lookup[i_zmn1], lookup[i], lookup[i_zpl1], lookup[i_zpl2]);

        return cubic_interpolation(w_rho, lookup_rmn1, lookup_r, lookup_rpl1, lookup_rpl2);
    }

    __global__ void coupling_kernel(const int *n, const float *m, const float *x_pos, const float *y_pos,
                                    const float *z_pos, const float *re_lookup_pl, const float *im_lookup_pl,
                                    const float *re_lookup_mn, const float *im_lookup_mn,
                                    const float *re_in_vec, const float *im_in_vec,
                                    float  *re_result, float  *im_result)
    {
        unsigned int i1 = blockIdx.x * blockDim.x + threadIdx.x;
        if(i1 >= NUMBER_OF_UNKNOWNS) return;

        const float x1 = x_pos[i1];
        const float y1 = y_pos[i1];
        const float z1 = z_pos[i1];

        const int n1 = n[i1];
        const float m1 = m[i1];

        re_result[i1] = 0.0;
        im_result[i1] = 0.0;

        for (int i2=0; i2<NUMBER_OF_UNKNOWNS; i2++)
        {
            float x21 = x1 - x_pos[i2];
            float y21 = y1 - y_pos[i2];
            float sz21 = z1 + z_pos[i2];
            float dz21 = z1 - z_pos[i2];

            const int n2 = n[i2];
            const float m2 = m[i2];

            float rho = sqrt(x21*x21+y21*y21);
            float phi = atan2(y21,x21);

            int rho_idx = (int) floor((rho - MIN_RHO) / LOOKUP_RESOLUTION);
            float rho_w = (rho - MIN_RHO) / LOOKUP_RESOLUTION - floor((rho - MIN_RHO) / LOOKUP_RESOLUTION);

            int sz_idx = (int) floor((sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION);
            float sz_w = (sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION - floor((sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION);

            int dz_idx = (int) floor((dz21 - MIN_Z_DIFF) / LOOKUP_RESOLUTION);
            float dz_w = (dz21 - MIN_Z_DIFF) / LOOKUP_RESOLUTION
                            - floor((dz21 - MIN_Z_DIFF) /  LOOKUP_RESOLUTION);

            float re_si_pl = lookup_2D(rho_idx, rho_w, sz_idx, sz_w, n1, n2, re_lookup_pl);
            float im_si_pl = lookup_2D(rho_idx, rho_w, sz_idx, sz_w, n1, n2, im_lookup_pl);
            float re_si_mn = lookup_2D(rho_idx, rho_w, dz_idx, dz_w, n1, n2, re_lookup_mn);
            float im_si_mn = lookup_2D(rho_idx, rho_w, dz_idx, dz_w, n1, n2, im_lookup_mn);

            float re_eimphi = cosf((m2 - m1) * phi);
            float im_eimphi = sinf((m2 - m1) * phi);

            float re_w = re_eimphi * (re_si_pl + re_si_mn) - im_eimphi * (im_si_pl + im_si_mn);
            float im_w = im_eimphi * (re_si_pl + re_si_mn) + re_eimphi * (im_si_pl + im_si_mn);

            re_result[i1] += re_w * re_in_vec[i2] - im_w * im_in_vec[i2];
            im_result[i1] += re_w * im_in_vec[i2] + im_w * re_in_vec[i2];
        }
    }"""

# This cuda kernel multiplies the coupling matrix to a vector. It is based on linear interpolation of the lookup
# table.
#
# input arguments of the cuda kernel:
# n (np.uint32):  n[i] contains the mutlipole multi-index with regard to self.l_max and self.m_max of the i-th
#                 entry of a system vector
# m (np.float32): m[i] contains the multipole order with regard to particle.l_max and particle.m_max
# x_pos (np.float32): x_pos[i] contains the respective particle x-position
# y_pos (np.float32): y_pos[i] contains the respective particle y-position
# re_lookup (np.float32): the real part of the lookup table, in the format [r, n1, n2]
# im_lookup (np.float32): the imaginary part of the lookup table, in the format [r, n1, n2]
# re_in_vec (np.float32): the real part of the vector to be multiplied with the coupling matrix
# im_in_vec (np.float32): the imaginary part of the vector to be multiplied with the coupling matrix
# re_result_vec (np.float32): the real part of the vector into which the result is written
# im_result_vec (np.float32): the imaginary part of the vector into which the result is written
linear_radial_lookup_source = """
    #define BLOCKSIZE %i
    #define NUMBER_OF_UNKNOWNS %i
    #define MIN_RHO %f
    #define LOOKUP_RESOLUTION %f

    __global__ void coupling_kernel(const int *n, const float *m, const float *x_pos, const float *y_pos,
                                    const float *re_lookup, const float *im_lookup, const float *re_in_vec,
                                    const float *im_in_vec, float  *re_result, float  *im_result)
    {
        unsigned int i1 = blockIdx.x * blockDim.x + threadIdx.x;
        if(i1 >= NUMBER_OF_UNKNOWNS) return;

        const float x1 = x_pos[i1];
        const float y1 = y_pos[i1];

        const int n1 = n[i1];
        const float m1 = m[i1];

        re_result[i1] = 0.0;
        im_result[i1] = 0.0;

        for (int i2=0; i2<NUMBER_OF_UNKNOWNS; i2++)
        {
            float x21 = x1 - x_pos[i2];
            float y21 = y1 - y_pos[i2];

            const int n2 = n[i2];
            const float m2 = m[i2];

            float r = sqrt(x21*x21+y21*y21);
            float phi = atan2(y21,x21);

            int r_idx = (int) floor((r - MIN_RHO) / LOOKUP_RESOLUTION);
            float w = (r - MIN_RHO) / LOOKUP_RESOLUTION - floor((r - MIN_RHO) / LOOKUP_RESOLUTION);

            int idx = r_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;
            int idx_pl_1 = (r_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            float re_si = re_lookup[idx] + w * (re_lookup[idx_pl_1] - re_lookup[idx]);
            float im_si = im_lookup[idx] + w * (im_lookup[idx_pl_1] - im_lookup[idx]);

            float re_eimphi = cosf((m2 - m1) * phi);
            float im_eimphi = sinf((m2 - m1) * phi);

            float re_w = re_eimphi * re_si - im_eimphi * im_si;
            float im_w = im_eimphi * re_si + re_eimphi * im_si;

            re_result[i1] += re_w * re_in_vec[i2] - im_w * im_in_vec[i2];
            im_result[i1] += re_w * im_in_vec[i2] + im_w * re_in_vec[i2];

        }
    }"""

# This cuda kernel multiplies the coupling matrix to a vector. It is based on cubic Hermite spline interpolation of the
# lookup table.
#
# input arguments of the cuda kernel:
# n (np.uint32):  n[i] contains the mutlipole multi-index with regard to self.l_max and self.m_max of the i-th
#                 entry of a system vector
# m (np.float32): m[i] contains the multipole order with regard to particle.l_max and particle.m_max
# x_pos (np.float32): x_pos[i] contains the respective particle x-position
# y_pos (np.float32): y_pos[i] contains the respective particle y-position
# re_lookup (np.float32): the real part of the lookup table, in the format [r, n1, n2]
# im_lookup (np.float32): the imaginary part of the lookup table, in the format [r, n1, n2]
# re_in_vec (np.float32): the real part of the vector to be multiplied with the coupling matrix
# im_in_vec (np.float32): the imaginary part of the vector to be multiplied with the coupling matrix
# re_result_vec (np.float32): the real part of the vector into which the result is written
# im_result_vec (np.float32): the imaginary part of the vector into which the result is written
cubic_radial_lookup_source = """
    #define BLOCKSIZE %i
    #define NUMBER_OF_UNKNOWNS %i
    #define MIN_RHO %f
    #define LOOKUP_RESOLUTION %f

    __device__ float cubic_interpolation(float const w, float const lookup_imn1, float const lookup_i,
                                         float const lookup_ipl1, float const lookup_ipl2)
    {
        return ((-w*w*w+2*w*w-w) * lookup_imn1 + (3*w*w*w-5*w*w+2) * lookup_i + (-3*w*w*w+4*w*w+w) * lookup_ipl1
                + (w*w*w-w*w) * lookup_ipl2) / 2;
    }

    __global__ void coupling_kernel(const int *n, const float *m, const float *x_pos, const float *y_pos,
                                    const float *re_lookup, const float *im_lookup, const float *re_in_vec,
                                    const float *im_in_vec, float  *re_result, float  *im_result)
    {
        unsigned int i1 = blockIdx.x * blockDim.x + threadIdx.x;
        if(i1 >= NUMBER_OF_UNKNOWNS) return;

        const float x1 = x_pos[i1];
        const float y1 = y_pos[i1];

        const int n1 = n[i1];
        const float m1 = m[i1];

        re_result[i1] = 0.0;
        im_result[i1] = 0.0;

        for (int i2=0; i2<NUMBER_OF_UNKNOWNS; i2++)
        {
            float x21 = x1 - x_pos[i2];
            float y21 = y1 - y_pos[i2];

            const int n2 = n[i2];
            const float m2 = m[i2];

            float r = sqrt(x21*x21+y21*y21);
            float phi = atan2(y21,x21);

            int r_idx = (int) floor((r - MIN_RHO) / LOOKUP_RESOLUTION);
            float w = (r - MIN_RHO) / LOOKUP_RESOLUTION - floor((r - MIN_RHO) / LOOKUP_RESOLUTION);

            int idx_mn1 = (r_idx - 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;
            int idx = r_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;
            int idx_pl1 = (r_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;
            int idx_pl2 = (r_idx + 2) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            float re_si = cubic_interpolation(w, re_lookup[idx_mn1], re_lookup[idx], re_lookup[idx_pl1],
                                              re_lookup[idx_pl2]);
            float im_si = cubic_interpolation(w, im_lookup[idx_mn1], im_lookup[idx], im_lookup[idx_pl1],
                                              im_lookup[idx_pl2]);

            float re_eimphi = cosf((m2 - m1) * phi);
            float im_eimphi = sinf((m2 - m1) * phi);

            float re_w = re_eimphi * re_si - im_eimphi * im_si;
            float im_w = im_eimphi * re_si + re_eimphi * im_si;

            re_result[i1] += re_w * re_in_vec[i2] - im_w * im_in_vec[i2];
            im_result[i1] += re_w * im_in_vec[i2] + im_w * re_in_vec[i2];
        }
    }"""

# This cuda kernel is used for the calculation of volume lookup tables (see smuthi.particle_coupling module).
volume_lookup_assembly_code = """
    #define BLOCKSIZE %i
    #define RHO_ARRAY_LENGTH %i
    #define Z_ARRAY_LENGTH %i
    #define K_ARRAY_LENGTH %i

    __global__ void helper(const float *re_bes_jac, const float *im_bes_jac, const float *re_belbee,
                            const float *im_belbee, const float *re_d_kappa, const float *im_d_kappa,
                            float *re_result, float  *im_result)
    {
        unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
        if(i >= RHO_ARRAY_LENGTH * Z_ARRAY_LENGTH) return;

        unsigned int i_rho = i / Z_ARRAY_LENGTH;
        unsigned int i_z = i %% Z_ARRAY_LENGTH;

        float re_res = 0.0;
        float im_res = 0.0;

        int i_kr = i_rho * K_ARRAY_LENGTH;
        int i_kz = i_z * K_ARRAY_LENGTH;

        float re_integrand_kp1 = re_bes_jac[i_kr] * re_belbee[i_kz] - im_bes_jac[i_kr] * im_belbee[i_kz];
        float im_integrand_kp1 = re_bes_jac[i_kr] * im_belbee[i_kz] + im_bes_jac[i_kr] * re_belbee[i_kz];

        for (int i_k=0; i_k<(K_ARRAY_LENGTH-1); i_k++)
        {
            i_kr = i_rho * K_ARRAY_LENGTH + i_k;
            i_kz = i_z * K_ARRAY_LENGTH + i_k;

            float re_integrand = re_integrand_kp1;
            float im_integrand = im_integrand_kp1;

            re_integrand_kp1 = re_bes_jac[i_kr+1] * re_belbee[i_kz+1] - im_bes_jac[i_kr+1] * im_belbee[i_kz+1];
            im_integrand_kp1 = re_bes_jac[i_kr+1] * im_belbee[i_kz+1] + im_bes_jac[i_kr+1] * re_belbee[i_kz+1];

            float re_sint = re_integrand + re_integrand_kp1;
            float im_sint = im_integrand + im_integrand_kp1;

            re_res += 0.5 * (re_sint * re_d_kappa[i_k] - im_sint * im_d_kappa[i_k]);
            im_res += 0.5 * (re_sint * im_d_kappa[i_k] + im_sint * re_d_kappa[i_k]);
        }

        re_result[i] = re_res;
        im_result[i] = im_res;

    }"""

# This cuda kernel is used for the calculation of radial lookup tables (see smuthi.particle_coupling module).
radial_lookup_assembly_code = """
    #define BLOCKSIZE %i
    #define RHO_ARRAY_LENGTH %i
    #define K_ARRAY_LENGTH %i

    __global__ void helper(const float *re_bes_jac, const float *im_bes_jac, const float *re_belbee,
                            const float *im_belbee, const float *re_d_kappa, const float *im_d_kappa,
                            float *re_result, float  *im_result)
    {
        unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
        if(i >= RHO_ARRAY_LENGTH) return;

        float re_res = 0.0;
        float im_res = 0.0;

        int i_kr = i * K_ARRAY_LENGTH;

        float re_integrand_kp1 = re_bes_jac[i_kr] * re_belbee[0] - im_bes_jac[i_kr] * im_belbee[0];
        float im_integrand_kp1 = re_bes_jac[i_kr] * im_belbee[0] + im_bes_jac[i_kr] * re_belbee[0];

        for (int i_k=0; i_k<(K_ARRAY_LENGTH-1); i_k++)
        {
            i_kr = i * K_ARRAY_LENGTH + i_k;

            float re_integrand = re_integrand_kp1;
            float im_integrand = im_integrand_kp1;

            re_integrand_kp1 = re_bes_jac[i_kr+1] * re_belbee[i_k+1] - im_bes_jac[i_kr+1] * im_belbee[i_k+1];
            im_integrand_kp1 = re_bes_jac[i_kr+1] * im_belbee[i_k+1] + im_bes_jac[i_kr+1] * re_belbee[i_k+1];

            float re_sint = re_integrand + re_integrand_kp1;
            float im_sint = im_integrand + im_integrand_kp1;

            re_res += 0.5 * (re_sint * re_d_kappa[i_k] - im_sint * im_d_kappa[i_k]);
            im_res += 0.5 * (re_sint * im_d_kappa[i_k] + im_sint * re_d_kappa[i_k]);
        }

        re_result[i] = re_res;
        im_result[i] = im_res;

    }"""

# This cuda kernel is used for the evaluation of the electric field of plane wave expansions.
pwe_electric_field_evaluation_code = """
    #define LEN_X %i
    #define LEN_K %i
    #define LEN_A %i
    #define RE_ONE_OVER_K %.10f
    #define IM_ONE_OVER_K %.10f

    __device__ void alpha_integral(const int i_k, const float x, const float y, const float z,
                                   const float re_kp, const float im_kp, const float re_kz, const float im_kz,
                                   const float *alpha_array,
                                   const float *re_g_te_array, const float *im_g_te_array,
                                   const float *re_g_tm_array, const float *im_g_tm_array,
                                   float *re_integral_x, float *im_integral_x,
                                   float *re_integral_y, float *im_integral_y,
                                   float *re_integral_z, float *im_integral_z)
    {
        float a = 0.0;

        float re_x_result = 0.0;
        float im_x_result = 0.0;
        float re_y_result = 0.0;
        float im_y_result = 0.0;
        float re_z_result = 0.0;
        float im_z_result = 0.0;

        float re_e_x_alpha_integrand = 0.0;
        float im_e_x_alpha_integrand = 0.0;
        float re_e_y_alpha_integrand = 0.0;
        float im_e_y_alpha_integrand = 0.0;
        float re_e_z_alpha_integrand = 0.0;
        float im_e_z_alpha_integrand = 0.0;

        for (int i_a=0; i_a<LEN_A; i_a++)
        {
            float a_old = a;
            a = alpha_array[i_a];
            float d_a = a - a_old;
            float sin_a = sinf(a);
            float cos_a = cosf(a);

            float re_ikr = -im_kp * (cos_a * x + sin_a * y) - im_kz * z;
            float im_ikr = re_kp * (cos_a * x + sin_a * y) + re_kz * z;

            float ereikr = expf(re_ikr);
            float re_eikr =  ereikr * cosf(im_ikr);
            float im_eikr =  ereikr * sinf(im_ikr);

            // integrand
            float re_e_x_alpha_integrand_old = re_e_x_alpha_integrand;
            float im_e_x_alpha_integrand_old = im_e_x_alpha_integrand;
            float re_e_y_alpha_integrand_old = re_e_y_alpha_integrand;
            float im_e_y_alpha_integrand_old = im_e_y_alpha_integrand;
            float re_e_z_alpha_integrand_old = re_e_z_alpha_integrand;
            float im_e_z_alpha_integrand_old = im_e_z_alpha_integrand;

            int i_ka = i_k * LEN_A + i_a;

            // pol=TE
            float re_g = re_g_te_array[i_ka];
            float im_g = im_g_te_array[i_ka];

            float re_geikr = re_g * re_eikr - im_g * im_eikr;
            float im_geikr = re_g * im_eikr + im_g * re_eikr;

            re_e_x_alpha_integrand = -sin_a * re_geikr;
            im_e_x_alpha_integrand = -sin_a * im_geikr;
            re_e_y_alpha_integrand = cos_a * re_geikr;
            im_e_y_alpha_integrand = cos_a * im_geikr;

            // pol=TM
            re_g = re_g_tm_array[i_ka];
            im_g = im_g_tm_array[i_ka];

            re_geikr = re_g * re_eikr - im_g * im_eikr;
            im_geikr = re_g * im_eikr + im_g * re_eikr;

            float re_kzgeikr = re_kz * re_geikr - im_kz * im_geikr;
            float im_kzgeikr = re_kz * im_geikr + im_kz * re_geikr;

            float re_kzkgeikr = re_kzgeikr * RE_ONE_OVER_K - im_kzgeikr * IM_ONE_OVER_K;
            float im_kzkgeikr = re_kzgeikr * IM_ONE_OVER_K + im_kzgeikr * RE_ONE_OVER_K;

            re_e_x_alpha_integrand += cos_a * re_kzkgeikr;
            im_e_x_alpha_integrand += cos_a * im_kzkgeikr;
            re_e_y_alpha_integrand += sin_a * re_kzkgeikr;
            im_e_y_alpha_integrand += sin_a * im_kzkgeikr;

            float re_kpgeikr = re_kp * re_geikr - im_kp * im_geikr;
            float im_kpgeikr = re_kp * im_geikr + im_kp * re_geikr;
            re_e_z_alpha_integrand = -(re_kpgeikr * RE_ONE_OVER_K - im_kpgeikr * IM_ONE_OVER_K);
            im_e_z_alpha_integrand = -(re_kpgeikr * IM_ONE_OVER_K + im_kpgeikr * RE_ONE_OVER_K);

            if (i_a>0)
            {
                re_x_result += 0.5 * d_a * (re_e_x_alpha_integrand + re_e_x_alpha_integrand_old);
                im_x_result += 0.5 * d_a * (im_e_x_alpha_integrand + im_e_x_alpha_integrand_old);
                re_y_result += 0.5 * d_a * (re_e_y_alpha_integrand + re_e_y_alpha_integrand_old);
                im_y_result += 0.5 * d_a * (im_e_y_alpha_integrand + im_e_y_alpha_integrand_old);
                re_z_result += 0.5 * d_a * (re_e_z_alpha_integrand + re_e_z_alpha_integrand_old);
                im_z_result += 0.5 * d_a * (im_e_z_alpha_integrand + im_e_z_alpha_integrand_old);
            }
        }
        re_integral_x[0] = re_x_result;
        im_integral_x[0] = im_x_result;
        re_integral_y[0] = re_y_result;
        im_integral_y[0] = im_y_result;
        re_integral_z[0] = re_z_result;
        im_integral_z[0] = im_z_result;
    }

    __global__ void electric_field(const float *re_kp_array, const float *im_kp_array,
                                   const float *re_kz_array, const float *im_kz_array,
                                   const float *alpha_array,
                                   const float *x_array, const float *y_array, const float *z_array,
                                   const float *re_g_te_array, const float *im_g_te_array,
                                   const float *re_g_tm_array, const float *im_g_tm_array,
                                   float *re_e_x, float *im_e_x, float *re_e_y, float *im_e_y,
                                   float *re_e_z, float *im_e_z)
    {
        unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= LEN_X) return;

        float x = x_array[i];
        float y = y_array[i];
        float z = z_array[i];

        float re_kp = 0.0;
        float im_kp = 0.0;

        float re_x_result = 0.0;
        float im_x_result = 0.0;
        float re_y_result = 0.0;
        float im_y_result = 0.0;
        float re_z_result = 0.0;
        float im_z_result = 0.0;

        float re_x_kappa_integrand = 0.0;
        float im_x_kappa_integrand = 0.0;
        float re_y_kappa_integrand = 0.0;
        float im_y_kappa_integrand = 0.0;
        float re_z_kappa_integrand = 0.0;
        float im_z_kappa_integrand = 0.0;

        for (int i_k=0; i_k<LEN_K; i_k++)
        {
            float re_kp_old = re_kp;
            float im_kp_old = im_kp;

            re_kp = re_kp_array[i_k];
            im_kp = im_kp_array[i_k];

            float re_d_kp = re_kp - re_kp_old;
            float im_d_kp = im_kp - im_kp_old;

            float re_kz = re_kz_array[i_k];
            float im_kz = im_kz_array[i_k];

            float re_x_kappa_integrand_old = re_x_kappa_integrand;
            float im_x_kappa_integrand_old = im_x_kappa_integrand;
            float re_y_kappa_integrand_old = re_y_kappa_integrand;
            float im_y_kappa_integrand_old = im_y_kappa_integrand;
            float re_z_kappa_integrand_old = re_z_kappa_integrand;
            float im_z_kappa_integrand_old = im_z_kappa_integrand;

            float re_x_alpha_itgrl = 0.0;
            float im_x_alpha_itgrl = 0.0;
            float re_y_alpha_itgrl = 0.0;
            float im_y_alpha_itgrl = 0.0;
            float re_z_alpha_itgrl = 0.0;
            float im_z_alpha_itgrl = 0.0;

            alpha_integral(i_k, x, y, z, re_kp, im_kp, re_kz, im_kz, alpha_array,
                           re_g_te_array, im_g_te_array, re_g_tm_array, im_g_tm_array,
                           &re_x_alpha_itgrl, &im_x_alpha_itgrl, &re_y_alpha_itgrl, &im_y_alpha_itgrl,
                           &re_z_alpha_itgrl, &im_z_alpha_itgrl);

            re_x_kappa_integrand = re_x_alpha_itgrl * re_kp - im_x_alpha_itgrl * im_kp;
            im_x_kappa_integrand = re_x_alpha_itgrl * im_kp + im_x_alpha_itgrl * re_kp;
            re_y_kappa_integrand = re_y_alpha_itgrl * re_kp - im_y_alpha_itgrl * im_kp;
            im_y_kappa_integrand = re_y_alpha_itgrl * im_kp + im_y_alpha_itgrl * re_kp;
            re_z_kappa_integrand = re_z_alpha_itgrl * re_kp - im_z_alpha_itgrl * im_kp;
            im_z_kappa_integrand = re_z_alpha_itgrl * im_kp + im_z_alpha_itgrl * re_kp;

            if (i_k>0)
            {
                re_x_result += 0.5 * (re_d_kp * (re_x_kappa_integrand + re_x_kappa_integrand_old)
                                      - im_d_kp * (im_x_kappa_integrand + im_x_kappa_integrand_old));
                im_x_result += 0.5 * (re_d_kp * (im_x_kappa_integrand + im_x_kappa_integrand_old)
                                      + im_d_kp * (re_x_kappa_integrand + re_x_kappa_integrand_old));
                re_y_result += 0.5 * (re_d_kp * (re_y_kappa_integrand + re_y_kappa_integrand_old)
                                      - im_d_kp * (im_y_kappa_integrand + im_y_kappa_integrand_old));
                im_y_result += 0.5 * (re_d_kp * (im_y_kappa_integrand + im_y_kappa_integrand_old)
                                      + im_d_kp * (re_y_kappa_integrand + re_y_kappa_integrand_old));
                re_z_result += 0.5 * (re_d_kp * (re_z_kappa_integrand + re_z_kappa_integrand_old)
                                      - im_d_kp * (im_z_kappa_integrand + im_z_kappa_integrand_old));
                im_z_result += 0.5 * (re_d_kp * (im_z_kappa_integrand + im_z_kappa_integrand_old)
                                      + im_d_kp * (re_z_kappa_integrand + re_z_kappa_integrand_old));
            }
        }
        re_e_x[i] = re_x_result;
        im_e_x[i] = im_x_result;
        re_e_y[i] = re_y_result;
        im_e_y[i] = im_y_result;
        re_e_z[i] = re_z_result;
        im_e_z[i] = im_z_result;
    }
"""


#!/usr/bin/env python2
from __future__ import print_function
import sys
sys.path.append('../lib')
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.path import Path
from matplotlib.patches import PathPatch
import string

import protocols
import model as m
from parameters import simvc, get_qc
from parameters import simvc_fix, simvc_fix_typical_values
from parameters import simvc_typical_values
from releakcorrect import I_releak, score_leak, protocol_leak_check

from scipy.optimize import fmin
# Set seed
np.random.seed(101)

savedir = './figs'
if not os.path.isdir(savedir):
    os.makedirs(savedir)

#refcell = 'D19'

def get_fix_param(var, val):
    """
        var: variable name.
        val: variable value to fix.
    """
    out = {}
    for i, j in zip(var, val):
        out[i] = j
    return out

def rmsd_compute(t1, t2):
    # Normalised RMSD value between trace 1 ``t1`` and trace 2 ``t2``
    #
    # Note, usually normalise to data, so
    # - ``t2`` data (or anything as reference)
    # - ``t1`` simulation (or anything for comparison)
    return np.sqrt(np.mean((t1 - t2) ** 2)) / np.sqrt(np.mean(t2 ** 2))

#
# Protocol info
#
protocol_funcs = {
    'staircaseramp': protocols.leak_staircase,
    'pharma': protocols.pharma,  # during drug application
    'apab': 'protocol-apab.csv',
    'apabv3': 'protocol-apabv3.csv',
    'ap05hz': 'protocol-ap05hz.csv',
    'ap1hz': 'protocol-ap1hz.csv',
    'ap2hz': 'protocol-ap2hz.csv',
    'sactiv': protocols.sactiv,
    'sinactiv': protocols.sinactiv,
}
protocol_dir = '../protocol-time-series'
protocol_list = [
    'staircaseramp',
    # 'sactiv',
    # 'sinactiv',
    'pharma',
    'apab',
    'apabv3',
    # 'ap05hz',
    'ap1hz',
    'ap2hz',
]
validation_idx = [
    None,
    # 1,
    # 2,
    3,
    4,
    5,
    # 6,
    7,
    8,
]

# IV protocol special treatment
protocol_iv = [
    'sactiv',
    'sinactiv',
]
protocol_iv_times = {
    'sactiv': protocols.sactiv_times,
    'sinactiv': protocols.sinactiv_times,
}
protocol_iv_convert = {
    'sactiv': protocols.sactiv_convert,
    'sinactiv': protocols.sinactiv_convert,
}
protocol_iv_args = {
    'sactiv': protocols.sactiv_iv_arg,
    'sinactiv': protocols.sinactiv_iv_arg,
}
protocol_iv_v = {
    'sactiv': protocols.sactiv_v,
    'sinactiv': protocols.sinactiv_v,
}

data_dir_staircase = '../data'
data_dir = '../data-autoLC'
file_dir = '../../hERGRapidCharacterisation/room-temperature-only/out'
file_dir2 = './out'
file_list = [
        'herg25oc1',
        ]
temperatures = np.array([25.0])
temperatures += 273.15  # in K
fit_seed = '542811797'
fit_seed2 = '717354021'

file_name = file_list[0]
temperature = temperatures[0]

# Load RMSD matrix
rmsd_matrix_file = '../../hERGRapidCharacterisation/room-temperature-only/figs/rmsd-hist-%s-autoLC-releak/rmsd-matrix.txt' \
                   % file_name
rmsd_cells_file = '../../hERGRapidCharacterisation/room-temperature-only/figs/rmsd-hist-%s-autoLC-releak/rmsd-matrix-cells.txt' \
                  % file_name

rmsd_matrix = np.loadtxt(rmsd_matrix_file)

with open(rmsd_matrix_file, 'r') as f:
    rmsd_prt = f.readline().strip('\n').strip('#').split()

rmsd_cells = []
with open(rmsd_cells_file, 'r') as f:
    for l in f:
        if not l.startswith('#'):
            rmsd_cells.append(l.strip('\n').split('-')[1])

rmsd_matrix_file2 = './out/rmsd-hist-%s-fixkinetics-simvclinleak-scheme3/rmsd-matrix.txt' \
                   % file_name
rmsd_cells_file2 = './out/rmsd-hist-%s-fixkinetics-simvclinleak-scheme3/rmsd-matrix-cells.txt' \
                  % file_name

rmsd_matrix2 = np.loadtxt(rmsd_matrix_file2)

with open(rmsd_matrix_file2, 'r') as f:
    rmsd_prt2 = f.readline().strip('\n').strip('#').split()

rmsd_cells2 = []
with open(rmsd_cells_file2, 'r') as f:
    for l in f:
        if not l.startswith('#'):
            rmsd_cells2.append(l.strip('\n').split('-')[1])

rankedlabels = [r'$*$',
                u'\u2021',
                r'#',
                u'\u2666']

#
# Do a very very tailored version........ :(
#
fig = plt.figure(figsize=(16, 15))
bigxgap = 12
n_xgrid = 84
bigygap = 5
n_ygrid = 31
grid = plt.GridSpec(2 * n_ygrid + 1 * bigygap, 3 * n_xgrid + 2 * bigxgap,
                    hspace=0.0, wspace=0.0)
axes = np.empty([10, int(len(protocol_list) / 2)], dtype=object)
# long list here:
for i in range(int(len(protocol_list) / 2)):
    i_grid = i * (n_xgrid + bigxgap)
    f_grid = (i + 1) * n_xgrid + i * bigxgap

    # First 'row'
    axes[0, i] = fig.add_subplot(grid[0:3, i_grid:f_grid])
    axes[0, i].set_xticklabels([])
    axes[1, i] = fig.add_subplot(grid[3:9, i_grid:f_grid])
    axes[1, i].set_xticklabels([])
    axes[2, i] = fig.add_subplot(grid[9:15, i_grid:f_grid])
    axes[2, i].set_xticklabels([])
    axes[3, i] = fig.add_subplot(grid[15:21, i_grid:f_grid])
    # Histogram
    axes[4, i] = fig.add_subplot(grid[24:31, i_grid:f_grid])

    # Second 'row'
    n_shift = n_ygrid + bigygap
    axes[5, i] = fig.add_subplot(grid[n_shift+0:n_shift+3, i_grid:f_grid])
    axes[5, i].set_xticklabels([])
    axes[6, i] = fig.add_subplot(grid[n_shift+3:n_shift+9, i_grid:f_grid])
    axes[6, i].set_xticklabels([])
    axes[7, i] = fig.add_subplot(grid[n_shift+9:n_shift+15, i_grid:f_grid])
    axes[7, i].set_xticklabels([])
    axes[8, i] = fig.add_subplot(grid[n_shift+15:n_shift+21, i_grid:f_grid])
    # Histogram
    axes[9, i] = fig.add_subplot(grid[n_shift+24:n_shift+31, i_grid:f_grid])

    # Set x-labels
    axes[3, i].set_xlabel('Time (s)', fontsize=14)
    axes[4, i].set_xlabel('RRMSE', fontsize=14)
    axes[8, i].set_xlabel('Time (s)', fontsize=14)
    axes[9, i].set_xlabel('RRMSE', fontsize=14)

# Set labels
axes[0, 0].set_ylabel('Voltage\n(mV)', fontsize=14)
axes[1, 0].set_ylabel(u'Best\n(*)', fontsize=14, color='#d95f02')
axes[2, 0].set_ylabel(u'Median\n(\u2021)', fontsize=14, color='#d95f02')
axes[3, 0].set_ylabel(u'90%ile\n(#)', fontsize=14, color='#d95f02')
axes[4, 0].set_ylabel('Frequency\n(N=%s)' % len(rmsd_cells), fontsize=14)
axes[5, 0].set_ylabel('Voltage\n(mV)', fontsize=14)
axes[6, 0].set_ylabel(u'Best\n(*)', fontsize=14, color='#d95f02')
axes[7, 0].set_ylabel(u'Median\n(\u2021)', fontsize=14, color='#d95f02')
axes[8, 0].set_ylabel(u'90%ile\n(#)', fontsize=14, color='#d95f02')
axes[9, 0].set_ylabel('Frequency\n(N=%s)' % len(rmsd_cells), fontsize=14)

axes[2, 0].text(-0.3, 0.5, 'Current (pA)', rotation=90, fontsize=18,
        transform=axes[2, 0].transAxes, ha='center', va='center')
axes[7, 0].text(-0.275, 0.5, 'Current (pA)', rotation=90, fontsize=18,
        transform=axes[7, 0].transAxes, ha='center', va='center')

#
# Model
#
prt2model = {}
prt2fixkineticsmodel = {}
for prt in protocol_list:

    protocol_def = protocol_funcs[prt]
    if type(protocol_def) is str:
        protocol_def = '%s/%s' % (protocol_dir, protocol_def)

    prt2model[prt] = m.Model('../mmt-model-files/ideal-ikr.mmt',
                        protocol_def=protocol_def,
                        temperature=temperature,  # K
                        transform=None,
                        useFilterCap=False)  # ignore capacitive spike

    prt2fixkineticsmodel[prt] = m.Model(
                '../mmt-model-files/simplified-voltage-clamp-ikr-linleak.mmt',
                protocol_def=protocol_def,
                temperature=temperature,  # K
                transform=None,
                useFilterCap=False)  # ignore capacitive spike

#
# Plot
#
mid = []
upper = []
lower = []
mid2 = []
upper2 = []
lower2 = []
for i_prt, prt in enumerate(protocol_list):

    # Calculate axis index
    ai, aj = 5 * int(i_prt / 3), i_prt % 3

    # Title
    if prt == 'staircaseramp':
        axes[ai, aj].set_title('Calibration', fontsize=16)
    else:
        axes[ai, aj].set_title('Validation %s' % validation_idx[i_prt],
                fontsize=16)

    # Add label!
    axes[ai, aj].text(-0.1, 1.4, string.ascii_uppercase[i_prt],
                      transform=axes[ai, aj].transAxes, size=20,
                      weight='bold')

    # Time point
    times = np.loadtxt('%s/%s-%s-times.csv' % (data_dir, file_name,
        prt), delimiter=',', skiprows=1) * 1e3  # s -> ms

    # Protocol
    model = prt2model[prt]
    modelfixkinetics = prt2fixkineticsmodel[prt]
    # Set which parameters to be inferred
    modelfixkinetics.set_parameters([
        'ikr.g',
        #'voltageclamp.rseries',
        'voltageclamp.voffset_eff',
        'voltageclamp.gLeak'])
    if prt not in protocol_iv:
        times_sim = np.copy(times)
        voltage = model.voltage(times_sim)
    else:
        times_sim = protocol_iv_times[prt](times[1] - times[0])
        voltage = model.voltage(times_sim)
        voltage, t = protocol_iv_convert[prt](voltage, times_sim)
        assert(np.mean(np.abs(t - times)) < 1e-8)
    axes[ai, aj].set_ylim((np.min(voltage) - 10, np.max(voltage) + 15))

    # Plot protocol
    if prt not in protocol_iv:
        axes[ai, aj].plot(times * 1e-3, voltage, c='#696969')
    else:
        # protocol
        for i in range(voltage.shape[1]):
            axes[ai, aj].plot(times * 1e-3, voltage[:, i], c='#696969')

    # Calculate ranking
    rmsd = rmsd_matrix[:, rmsd_prt.index(prt)]
    best_cell = np.argmin(rmsd)
    median_cell = np.argsort(rmsd)[len(rmsd)//2]
    p90_cell = np.argsort(rmsd)[int(len(rmsd)*0.9)]
    rankedcells = [rmsd_cells[best_cell],
                   rmsd_cells[median_cell],
                   rmsd_cells[p90_cell]]
    rankedvalues = [rmsd[best_cell],
                    rmsd[median_cell],
                    rmsd[p90_cell]]
                    #rmsd[rmsd_cells.index(refcell)]]

    rmsd2 = rmsd_matrix2[:, rmsd_prt2.index(prt)]
    #best_cell2 = np.argmin(rmsd2)
    #median_cell2 = np.argsort(rmsd2)[len(rmsd2)//2]
    #p90_cell2 = np.argsort(rmsd2)[int(len(rmsd2)*0.9)]
    #NOTE Compare with the 'red' cells; not its own ranking!!
    best_cell2 = rmsd_cells2.index(rmsd_cells[best_cell])
    median_cell2 = rmsd_cells2.index(rmsd_cells[median_cell])
    p90_cell2 = rmsd_cells2.index(rmsd_cells[p90_cell])
    rankedcells2 = [rmsd_cells2[best_cell2],
                   rmsd_cells2[median_cell2],
                   rmsd_cells2[p90_cell2]]
    rankedvalues2 = [rmsd2[best_cell2],
                    rmsd2[median_cell2],
                    rmsd2[p90_cell2]]
                    #rmsd2[rmsd_cells2.index(refcell)]]

    # Parameters
    fn = '%s/%s-scheme3-simvclinleak/%s-cells-%s.txt' % \
         (file_dir2, file_name, file_name, fit_seed2)
    scheme3_cell_list = []
    with open(fn, 'r') as f:
        for l in f:
            if not l.startswith('#'):
                scheme3_cell_list.append(l.split()[0])

    param_file = '%s/%s-scheme3-simvclinleak/%s-solution_i-%s.txt' % \
            (file_dir2, file_name, file_name, fit_seed2)
    obtained_parameters_all = np.loadtxt(param_file)

    ikr_param = [
            'ikr.p1', 'ikr.p2', 'ikr.p3', 'ikr.p4',
            'ikr.p5', 'ikr.p6', 'ikr.p7', 'ikr.p8',
            ]
    p_ikr = np.loadtxt('%s/%s-scheme3-simvclinleak/%s-solution-%s.txt' % \
            (file_dir2, file_name, file_name, fit_seed2))

    for i_cell, cell in enumerate(rankedcells):
        # Data
        if prt == 'staircaseramp':
            data = np.loadtxt('%s/%s-%s-%s.csv' % (data_dir_staircase,
                    file_name, prt, cell), delimiter=',', skiprows=1)
        elif prt not in protocol_iv:
            data = np.loadtxt('%s/%s-%s-%s.csv' % (data_dir, file_name,
                    prt, cell), delimiter=',', skiprows=1)
            # Re-leak correct the leak corrected data...
            g_releak = fmin(score_leak, [0.0], args=(data, voltage, times,
                                protocol_leak_check[prt]), disp=False)
            data = I_releak(g_releak[0], data, voltage)
        else:
            data = np.loadtxt('%s/%s-%s-%s.csv' % (data_dir, file_name,
                    prt, cell), delimiter=',', skiprows=1)
            # Re-leak correct the leak corrected data...
            for i in range(data.shape[1]):
                g_releak = fmin(score_leak, [0.0], args=(data[:, i],
                                    voltage[:, i], times,
                                    protocol_leak_check[prt]), disp=False)
                data[:, i] = I_releak(g_releak[0], data[:, i], voltage[:, i])
        assert(len(data) == len(times))

        # Fitted parameters
        param_file = '%s/%s/%s-staircaseramp-%s-solution-%s.txt' % \
                (file_dir, file_name, file_name, cell, fit_seed)
        obtained_parameters = np.loadtxt(param_file) * 1e-3  # V, s -> mV, ms

        # For fix kinetics model
        rseal, cm, rseries = get_qc('../qc', file_name, cell)
        #print('Est. Rseal, Cm, Rseries:', rseal, cm, rseries, '(GOhm, pF, GOhm)')
        alpha = 0.8  # rseries %compensation
        simvc_fix_values = [cm, rseries * alpha, rseries]
        extra_fix = ['voltageclamp.rseries']
        updateELeakCorrection = False
        if updateELeakCorrection:
            leakbeforeparam = np.loadtxt('../qc/' + file_name + '-staircaseramp-leak_before.txt')
            leakafterparam = np.loadtxt('../qc/' + file_name + '-staircaseramp-leak_after.txt')
            cell_id_file = '../qc/%s-staircaseramp-cell_id.txt' % file_name
            cell_ids = []
            with open(cell_id_file, 'r') as f:
                for l in f:
                    if not l.startswith('#'):
                        cell_ids.append(l.split()[0])
            cell_idx = cell_ids.index(cell)
            ga, Ea = leakbeforeparam[cell_idx]
            gb, Eb = leakafterparam[cell_idx]
            ELeakCorrection = - (ga * Ea - gb * Eb) / (gb - ga)
            #print('E_Leak correction: ', ELeakCorrection, ' (mV)')
            if np.abs(ELeakCorrection) > 200: print('==' * 30, ga, Ea, gb, Eb)
            extra_fix += ['voltageclamp.ELeak']
            simvc_fix_values += [ELeakCorrection]
        fix_p = get_fix_param(ikr_param + simvc_fix + extra_fix,
                np.append(p_ikr, simvc_fix_values))
        modelfixkinetics.set_fix_parameters(fix_p)

        scheme3_cell_idx = scheme3_cell_list.index(cell)
        obtained_parameters2 = obtained_parameters_all[scheme3_cell_idx]

        # Simulation
        simulation = model.simulate(obtained_parameters, times_sim)
        simulationfixkinetics = modelfixkinetics.simulate(obtained_parameters2, times_sim)
        if prt != 'staircaseramp' and prt not in protocol_iv:
            # Re-leak correct the leak corrected simulationfixkinetics... TODO?
            g_releak_simulationfixkinetics = fmin(score_leak, [0.1], args=(simulationfixkinetics, voltage, times,
                                protocol_leak_check[prt]), disp=False)
            simulationfixkinetics = I_releak(g_releak_simulationfixkinetics[0], simulationfixkinetics, voltage)
        if prt in protocol_iv:
            simulation, t = protocol_iv_convert[prt](simulation, times_sim)
            assert(np.mean(np.abs(t - times)) < 1e-6)
            simulationfixkinetics, t = protocol_iv_convert[prt](
                    simulationfixkinetics, times_sim)
            assert(np.mean(np.abs(t - times)) < 1e-6)
            # Re-leak correct the leak corrected simulationfixkinetics... TODO?
            for i in range(simulationfixkinetics.shape[1]):
                g_releak_simulationfixkinetics = fmin(score_leak, [0.1], args=(simulationfixkinetics[:, i],
                                    voltage[:, i], times,
                                    protocol_leak_check[prt]), disp=False)
                simulationfixkinetics[:, i] = I_releak(g_releak_simulationfixkinetics[0], simulationfixkinetics[:, i], voltage[:, i])

        # Work out ylim
        maximum = np.percentile(simulation, 100)
        minimum = np.percentile(simulation, 0.0)
        amplitude = maximum - minimum
        if prt in ['apabv3', 'ap05hz']:
            maximum += 0.6 * amplitude
            minimum -= 0.6 * amplitude
        elif prt in ['apab', 'ap1hz']:
            maximum += 0.3 * amplitude
            minimum -= 0.3 * amplitude
        else:
            maximum += 0.15 * amplitude
            minimum -= 0.15 * amplitude
        axes[ai + i_cell + 1, aj].set_ylim([minimum, maximum])

        # Plot
        if prt not in protocol_iv:
            # recording
            axes[ai + i_cell + 1, aj].plot(times * 1e-3, data, lw=1, alpha=0.5,
                    c='#9ecae1', label='data')
            # simulation
            if prt == 'staircaseramp':
                axes[ai + i_cell + 1, aj].plot(times * 1e-3, simulation, lw=2,
                        c='#d95f02', label='model fit to data')
                axes[ai + i_cell + 1, aj].plot(times * 1e-3, simulationfixkinetics, lw=2,
                        c='#1b9e77', label='model (fix kinetics) fit to data')
            else:
                axes[ai + i_cell + 1, aj].plot(times * 1e-3, simulation, lw=2,
                        c='#d95f02', label='model prediction')
                axes[ai + i_cell + 1, aj].plot(times * 1e-3, simulationfixkinetics, lw=2,
                        c='#1b9e77', label='model (fix kinetics) prediction')
        else:
            iv_v = protocol_iv_v[prt]()  # mV
            # recording
            iv_i = protocols.get_corrected_iv(data, times,
                                              *protocol_iv_args[prt]())
            axes[ai + i_cell + 1, aj].plot(iv_v, iv_i / np.max(iv_i), lw=2,
                    alpha=0.25, c='#9ecae1', label='data')
            # simulation
            iv_i = protocols.get_corrected_iv(simulation, times,
                                              *protocol_iv_args[prt]())
            axes[ai + i_cell + 1, aj].plot(iv_v, iv_i / np.max(iv_i), lw=2,
                    alpha=1, c='#d95f02', label='model prediction')
            # simulationfixkinetics
            iv_i_fixkinetics = protocols.get_corrected_iv(simulationfixkinetics, times,
                                              *protocol_iv_args[prt]())
            axes[ai + i_cell + 1, aj].plot(iv_v, iv_i_fixkinetics / np.max(iv_i_fixkinetics), lw=2,
                    alpha=1, c='#1b9e77', label='model (fix kinetics) prediction')
            if prt == 'sactiv':
                axes[ai + i_cell + 1, aj].set_ylim([-0.05, 1.05])
            elif prt == 'sinactiv':
                axes[ai + i_cell + 1, aj].set_ylim([-5, 1.05])

        if False:
            print(prt, i_cell, cell)
            print('red', rmsd_compute(simulation, data))
            print('green', rmsd_compute(simulationfixkinetics, data))

    # Plot rmsd histogram
    rmse_min = min(np.min(rmsd), np.min(rmsd2))
    rmse_max = max(np.max(rmsd), np.max(rmsd2))
    rmse_range = rmse_max - rmse_min
    bins = np.linspace(rmse_min - 0.1 * rmse_range,
            rmse_max + 0.1 * rmse_range, 20)
    n, b, _ = axes[ai + 4, aj].hist(rmsd, bins=bins, color='#d95f02', alpha=0.25)
    n2, b2, _ = axes[ai + 4, aj].hist(rmsd2, bins=bins, color='#2ca02c', alpha=0.25)

    mid.append(np.percentile(rmsd, 50))
    upper.append(np.percentile(rmsd, 90))
    lower.append(np.percentile(rmsd, 10))
    mid2.append(np.percentile(rmsd2, 50))
    upper2.append(np.percentile(rmsd2, 90))
    lower2.append(np.percentile(rmsd2, 10))

    # Add labels
    rankedidx = []
    for i, v in enumerate(rankedvalues):
        idx = np.where(b <= v)[0][-1]
        if idx in rankedidx:
            print('Ref. marker might clash with other markers...')
            shift = 4
        else:
            shift = 0
        axes[ai + 4, aj].text((b[idx] + b[idx + 1]) / 2., n[idx] + 3 + shift,
                rankedlabels[i], fontsize=16, color='#d95f02',
                ha='center', va='center')
        if n[idx] > 0.8 * max(np.max(n), np.max(n2)):
            axes[ai + 4, aj].set_ylim([0, max(np.max(n2), np.max(n)) + 8 + shift])
        rankedidx.append(idx)

    rankedidx2 = []
    for i, v in enumerate(rankedvalues2):
        idx = np.where(b2 <= v)[0][-1]
        if idx in rankedidx2:
            print('Ref. marker might clash with other markers...')
            shift = 4
        elif idx in rankedidx:
            diff = np.abs(n[idx] - n2[idx])
            if diff < max(np.max(n2), np.max(n)) * 0.1:
                shift = max(np.max(n2), np.max(n)) * 0.2
            else:
                shift = 0
        else:
            shift = 0
        axes[ai + 4, aj].text((b2[idx] + b2[idx + 1]) / 2., n2[idx] + 3 + shift,
                rankedlabels[i], fontsize=16, color='#2ca02c',
                ha='center', va='center')
        if n2[idx] > 0.8 * max(np.max(n2), np.max(n)):
            axes[ai + 4, aj].set_ylim([0, max(np.max(n2), np.max(n)) + 8 + shift])
        rankedidx2.append(idx)

#
# Final adjustment and save
#
#axes[1, 0].legend()
#axes[1, 1].legend()
import matplotlib.patches as mpatches
data_patch = mpatches.Patch(color='#9ecae1', label='Data')
h1_patch = mpatches.Patch(color='#d95f02', label='Hypothesis 1: independent kinetics models')
h2_patch = mpatches.Patch(color='#1b9e77', label='Hypothesis 2: identical kinetics models')
axes[0, 0].legend(handles=[data_patch, h1_patch, h2_patch], loc='upper left', bbox_to_anchor=(-.025, 2.75), fontsize=14, ncol=3, columnspacing=5.5)
#grid.tight_layout(fig, pad=0.6, rect=(0.02, 0.0, 1, 0.99))  # not working...
#grid.update(wspace=0.2, hspace=0.0)
plt.savefig('%s/rmsd-hist-fix-kinetics-simvclinleak-scheme3-part1.pdf' % (savedir), bbox_inch='tight',
        pad_inches=0, format='pdf')
plt.savefig('%s/rmsd-hist-fix-kinetics-simvclinleak-scheme3-part1.png' % (savedir), bbox_inch='tight',
        pad_inches=0, dpi=300)

print('Done')

#
# Table
#
tex = ''
tex += '\\begin{tabularx}{\\textwidth}{@{}l' \
    + 'XXc' * (len(protocol_list) - 1) + 'XX@{}}\n'
tex += '\\toprule\n'
tex += '   '
for i_prt, prt in enumerate(protocol_list):
    # span 2 columns
    if prt == 'staircaseramp':
        tex += ' & \multicolumn{2}{c}{Cal.}'
    else:
        tex += ' & \multicolumn{2}{c}{Val.~%s}' % validation_idx[i_prt]

    if i_prt < len(protocol_list) - 1:
        tex += ' & \phantom{}'
tex += ' \\\\\n'

ii = 1
for i_prt, prt in enumerate(protocol_list):
    ii += 1
    tex += '\\cmidrule{%s-%s}' % (ii, ii + 1)
    ii += 2
tex += '\n'

tex += '   '
for i_prt, prt in enumerate(protocol_list):
    tex += ' & H1 & H2'
    if i_prt < len(protocol_list) - 1:
        tex += ' &'
tex += ' \\\\\n'

tex += '\\midrule\n'

tex += 'Median'
for i_prt, prt in enumerate(protocol_list):
    tex += ' & %.2f & %.2f' % (mid[i_prt], mid2[i_prt])
    if i_prt < len(protocol_list) - 1:
        tex += ' &'
tex += ' \\\\\n'

tex += '10\\textsuperscript{th} \\%ile'
for i_prt, prt in enumerate(protocol_list):
    tex += ' & %.2f & %.2f' % (lower[i_prt], lower2[i_prt])
    if i_prt < len(protocol_list) - 1:
        tex += ' &'
tex += ' \\\\\n'

tex += '90\\textsuperscript{th} \\%ile'
for i_prt, prt in enumerate(protocol_list):
    tex += ' & %.2f & %.2f' % (upper[i_prt], upper2[i_prt])
    if i_prt < len(protocol_list) - 1:
        tex += ' &'
tex += ' \\\\\n'
tex += '\\bottomrule\n'
tex += '\\end{tabularx}'

print(tex)


# Licensed under a 3-clause BSD style license - see LICENSE.rst
# This old version has been retrieved to use the Spectrum class to do photometry
# Functionalities have been removed

import abc
import math
from copy import deepcopy
from collections import OrderedDict

import numpy as np

from astropy.utils import lazyproperty
from astropy.io import ascii
import astropy.units as u
import astropy.constants as const
from astropy import cosmology

# from ._registry import Registry

__all__ = [
    "get_bandpass",
    "get_magsystem",
    "read_bandpass",
    "Bandpass",
    "Spectrum",
    "MagSystem",
    "SpectralMagSystem",
    "ABMagSystem",
    "CompositeMagSystem",
]

HC_ERG_AA = const.h.cgs.value * const.c.to(u.AA / u.s).value

# _BANDPASSES = Registry()
# _MAGSYSTEMS = Registry()

def get_bandpass(name):
    """Get a Bandpass from the registry by name."""
    if isinstance(name, Bandpass):
        return name
    return _BANDPASSES.retrieve(name)

def get_magsystem(name):
    """Get a MagSystem from the registry by name."""
    if isinstance(name, MagSystem):
        return name
    return _MAGSYSTEMS.retrieve(name)

def read_bandpass(
    fname, fmt="ascii", wave_unit=u.AA, trans_unit=u.dimensionless_unscaled, name=None
):
    """Read bandpass from two-column ASCII file containing wavelength and
    transmission in each line.

    Parameters
    ----------
    fname : str
        File name.
    fmt : {'ascii'}
        File format of file. Currently only ASCII file supported.
    wave_unit : `~astropy.units.Unit` or str, optional
        Wavelength unit. Default is Angstroms.
    trans_unit : `~astropy.units.Unit`, optional
        Transmission unit. Can be `~astropy.units.dimensionless_unscaled`,
        indicating a ratio of transmitted to incident photons, or units
        proportional to inverse energy, indicating a ratio of transmitted
        photons to incident energy. Default is ratio of transmitted to
        incident photons.
    name : str, optional
        Identifier. Default is `None`.

    Returns
    -------
    band : `~sncosmo.Bandpass`
    """

    if fmt != "ascii":
        raise ValueError(
            "format {0} not supported. Supported formats: 'ascii'".format(fmt)
        )
    t = ascii.read(fname, names=["wave", "trans"])
    return Bandpass(
        t["wave"], t["trans"], wave_unit=wave_unit, trans_unit=trans_unit, name=name
    )

class Bandpass(object):
    """Transmission as a function of spectral wavelength.

    Parameters
    ----------
    wave : list_like
        Wavelength. Monotonically increasing values.
    trans : list_like
        Transmission fraction.
    wave_unit : `~astropy.units.Unit` or str, optional
        Wavelength unit. Default is Angstroms.
    trans_unit : `~astropy.units.Unit`, optional
        Transmission unit. Can be `~astropy.units.dimensionless_unscaled`,
        indicating a ratio of transmitted to incident photons, or units
        proportional to inverse energy, indicating a ratio of transmitted
        photons to incident energy. Default is ratio of transmitted to
        incident photons.
    name : str, optional
        Identifier. Default is `None`.

    Examples
    --------
    >>> b = Bandpass([4000., 4200., 4400.], [0.5, 1.0, 0.5])
    >>> b.wave
    array([ 4000.,  4200.,  4400.])
    >>> b.trans
    array([ 0.5,  1. ,  0.5])
    >>> b.dwave
    array([ 200.,  200.,  200.])
    >>> b.wave_eff
    4200.0

    """

    def __init__(
        self,
        wave,
        trans,
        wave_unit=u.AA,
        trans_unit=u.dimensionless_unscaled,
        name=None,
    ):
        wave = np.asarray(wave, dtype=np.float64)
        trans = np.asarray(trans, dtype=np.float64)
        if wave.shape != trans.shape:
            raise ValueError("shape of wave and trans must match")
        if wave.ndim != 1:
            raise ValueError("only 1-d arrays supported")

        # Ensure that units are actually units and not quantities, so that
        # `to` method returns a float and not a Quantity.
        wave_unit = u.Unit(wave_unit)
        trans_unit = u.Unit(trans_unit)

        if wave_unit != u.AA:
            wave = wave_unit.to(u.AA, wave, u.spectral())

        # If transmission is in units of inverse energy, convert to
        # unitless transmission:
        #
        # (transmitted photons / incident photons) =
        #      (photon energy) * (transmitted photons / incident energy)
        #
        # where photon energy = h * c / lambda
        if trans_unit != u.dimensionless_unscaled:
            trans = (HC_ERG_AA / wave) * trans_unit.to(u.erg ** -1, trans)

        # Check that values are monotonically increasing.
        # We could sort them, but if this happens, it is more likely a user
        # error or faulty bandpass definition. So we leave it to the user to
        # sort them.
        if not np.all(np.ediff1d(wave) > 0.0):
            raise ValueError(
                "bandpass wavelength values must be monotonically"
                " increasing when supplied in wavelength or "
                "decreasing when supplied in energy/frequency."
            )
        self.wave = wave
        self._dwave = np.gradient(wave)
        self.trans = trans
        self.name = name

    @property
    def dwave(self):
        """Gradient of wavelengths, numpy.gradient(wave)."""
        return self._dwave

    @lazyproperty
    def wave_eff(self):
        """Effective wavelength of bandpass in Angstroms."""
        weights = self.trans * np.gradient(self.wave)
        return np.sum(self.wave * weights) / np.sum(weights)

    def to_unit(self, unit):
        """Return wavelength and transmission in new wavelength units.

        If the requested units are the same as the current units, self is
        returned.

        Parameters
        ----------
        unit : `~astropy.units.Unit` or str
            Target wavelength unit.

        Returns
        -------
        wave : `~numpy.ndarray`
        trans : `~numpy.ndarray`
        """

        if unit is u.AA:
            return self.wave, self.trans

        d = u.AA.to(unit, self.wave, u.spectral())
        t = self.trans
        if d[0] > d[-1]:
            d = np.flipud(d)
            t = np.flipud(t)
        return d, t

    def __repr__(self):
        name = ""
        if self.name is not None:
            name = " {0!r:s}".format(self.name)
        return "<Bandpass{0:s} at 0x{1:x}>".format(name, id(self))

class Spectrum(object):
    """A spectrum, representing wavelength and spectral density values.

    Parameters
    ----------
    wave : list_like
        Wavelength values.
    flux : list_like
        Spectral flux density values.
    error : list_like, optional
        1 standard deviation uncertainty on flux density values.
    wave_unit : `~astropy.units.Unit`
        Units.
    unit : `~astropy.units.BaseUnit`
        For now, only units with flux density in energy (not photon counts).
    z : float, optional
        Redshift of spectrum (default is `None`)
    dist : float, optional
        Luminosity distance in Mpc, used to adjust flux upon redshifting.
        The default is ``None``.
    meta : OrderedDict, optional
        Metadata.
    """

    def __init__(
        self,
        wave,
        flux,
        error=None,
        unit=(u.erg / u.s / u.cm ** 2 / u.AA),
        wave_unit=u.AA,
        z=None,
        dist=None,
        meta=None,
    ):
        self._wave = np.asarray(wave)
        self._flux = np.asarray(flux)
        self._wunit = wave_unit
        self._unit = unit
        self._z = z
        self._dist = dist

        if error is not None:
            self._error = np.asarray(error)
            if self._wave.shape != self._error.shape:
                raise ValueError("shape of wavelength and variance must match")
        else:
            self._error = None

        if meta is None:
            self.meta = OrderedDict()
        else:
            self.meta = deepcopy(meta)

        if self._wave.shape != self._flux.shape:
            raise ValueError("shape of wavelength and flux must match")
        if self._wave.ndim != 1:
            raise ValueError("only 1-d arrays supported")

    @property
    def wave(self):
        """Wavelength values."""
        return self._wave

    @property
    def flux(self):
        """Spectral flux density values"""
        return self._flux

    @property
    def error(self):
        """Uncertainty on flux density."""
        return self._error

    @property
    def wave_unit(self):
        """Units of wavelength."""
        return self._wunit

    @property
    def unit(self):
        """Units of flux density."""
        return self._unit

    @property
    def z(self):
        """Redshift of spectrum."""
        return self._z

    @z.setter
    def z(self, value):
        self._z = value

    @property
    def dist(self):
        """Distance to object in Mpc."""
        return self._dist

    @dist.setter
    def dist(self, value):
        self._dist = value

    def bandflux(self, band):
        """Perform synthentic photometry in a given bandpass.

        The bandpass transmission is interpolated onto the wavelength grid
        of the spectrum. The result is a weighted sum of the spectral flux
        density values (weighted by transmission values).

        Parameters
        ----------
        band : Bandpass object or name of registered bandpass.

        Returns
        -------
        bandflux : float
            Total flux in ph/s/cm^2. If part of bandpass falls
            outside the spectrum, `None` is returned instead.
        bandfluxerr : float
            Error on flux. Only returned if the `error` attribute is not
            `None`.
        """

        band = get_bandpass(band)
        bwave, btrans = band.to_unit(self._wunit)

        if bwave[0] < self._wave[0] or bwave[-1] > self._wave[-1]:
            return None

        mask = (self._wave > bwave[0]) & (self._wave < bwave[-1])
        d = self._wave[mask]
        f = self._flux[mask]

        # First convert to ergs/s/cm^2/(wavelength unit)...
        target_unit = u.erg / u.s / u.cm ** 2 / self._wunit
        if self._unit != target_unit:
            f = self._unit.to(target_unit, f, u.spectral_density(self._wunit, d))

        # Then convert ergs to photons: photons = Energy / (h * nu).
        f = f / const.h.cgs.value / self._wunit.to(u.Hz, d, u.spectral())

        trans = np.interp(d, bwave, btrans)
        binw = np.gradient(d)
        ftot = np.sum(f * trans * binw)

        if self._error is None:
            return ftot

        else:
            e = self._error[mask]

            # Do the same conversion as above
            if self._unit != target_unit:
                e = self._unit.to(target_unit, e, u.spectral_density(self._wunit, d))
            e = e / const.h.cgs.value / self._wunit.to(u.Hz, d, u.spectral())
            etot = np.sqrt(np.sum((e * binw) ** 2 * trans))
            return ftot, etot

    def redshifted_to(self, z, adjust_flux=False, dist=None, cosmo=None):
        """Return a new Spectrum object at a new redshift.

        The current redshift must be defined (self.z cannot be `None`).
        A factor of (1 + z) / (1 + self.z) is applied to the wavelength.
        The inverse factor is applied to the flux so that the bolometric
        flux (e.g., erg/s/cm^2) remains constant.

        .. note:: Currently this only works for units in ergs.

        Parameters
        ----------
        z : float
            Target redshift.
        adjust_flux : bool, optional
            If True, the bolometric flux is adjusted by
            ``F_out = F_in * (D_in / D_out) ** 2``, where ``D_in`` and
            ``D_out`` are current and target luminosity distances,
            respectively. ``D_in`` is self.dist. If self.dist is ``None``,
            the distance is calculated from the current redshift and
            given cosmology.
        dist : float, optional
            Output distance in Mpc. Used to adjust bolometric flux if
            ``adjust_flux`` is ``True``. Default is ``None`` which means
            that the distance is calculated from the redshift and the
            cosmology.
        cosmo : `~astropy.cosmology.Cosmology` instance, optional
            The cosmology used to estimate distances if dist is not given.
            Default is ``None``, which results in using the default
            cosmology.

        Returns
        -------
        spec : Spectrum object
            A new spectrum object at redshift z.
        """

        if self._z is None:
            raise ValueError(
                "Must set current redshift in order to redshift" " spectrum"
            )

        if self._wunit.physical_type == u.m.physical_type:
            factor = (1.0 + z) / (1.0 + self._z)
        elif self._wunit.physical_type == u.Hz.physical_type:
            factor = (1.0 + self._z) / (1.0 + z)
        else:
            raise ValueError("wavelength must be in wavelength or frequency")

        d = self._wave * factor
        f = self._flux / factor
        if self._error is not None:
            e = self._error / factor
        else:
            e = None

        if adjust_flux:
            if self._dist is None and self._z == 0.0:
                raise ValueError(
                    "When current redshift is 0 and adjust_flux "
                    "is requested, current distance must be "
                    "defined"
                )
            if dist is None and z == 0.0:
                raise ValueError(
                    "When redshift is 0 and adjust_flux "
                    "is requested, dist must be defined"
                )
            if cosmo is None:
                cosmo = cosmology.get_current()

            if self._dist is None:
                dist_in = cosmo.luminosity_distance(self._z)
            else:
                dist_in = self._dist

            if dist is None:
                dist = cosmo.luminosity_distance(z)

            if dist_in <= 0.0 or dist <= 0.0:
                raise ValueError("Distances must be greater than 0.")

            # Adjust the flux
            factor = (dist_in / dist) ** 2
            f *= factor
            if e is not None:
                e *= factor

        return Spectrum(
            d,
            f,
            error=e,
            z=z,
            dist=dist,
            meta=self.meta,
            unit=self._unit,
            wave_unit=self._wunit,
        )

class MagSystem(object):
    """An abstract base class for magnitude systems."""

    __metaclass__ = abc.ABCMeta

    def __init__(self, name=None):
        self._zpbandflux = {}
        self._name = name

    @abc.abstractmethod
    def _refspectrum_bandflux(self, band):
        """Flux of the fundamental spectrophotometric standard."""
        pass

    @property
    def name(self):
        """Name of magnitude system."""
        return self._name

    @name.setter
    def name(self, value):
        self._name = value

    def zpbandflux(self, band):
        """Flux of an object with magnitude zero in the given bandpass.

        Parameters
        ----------
        bandpass : `~sncosmo.spectral.Bandpass` or str

        Returns
        -------
        bandflux : float
            Flux in photons / s / cm^2.
        """

        band = get_bandpass(band)
        try:
            return self._zpbandflux[band]
        except KeyError:
            bandflux = self._refspectrum_bandflux(band)
            self._zpbandflux[band] = bandflux

        return bandflux

    def band_flux_to_mag(self, flux, band):
        """Convert flux (photons / s / cm^2) to magnitude."""
        return -2.5 * math.log10(flux / self.zpbandflux(band))

    def band_mag_to_flux(self, mag, band):
        """Convert magnitude to flux in photons / s / cm^2"""
        return self.zpbandflux(band) * 10.0 ** (-0.4 * mag)

class CompositeMagSystem(MagSystem):
    """A magnitude system defined in a specific set of bands.

    In each band, there is a fundamental standard with a known
    (generally non-zero) magnitude.

    Parameters
    ----------
    bands: iterable of `~sncosmo.Bandpass` or str
        The filters in the magnitude system.
    standards: iterable of `~sncosmo.MagSystem` or str,
        The spectrophotmetric flux standards for each band, in the
        same order as `bands`.
    offsets: list_like
        The magnitude of standard in the given band.
    """

    def __init__(self, bands, standards, offsets, name=None):
        super(CompositeMagSystem, self).__init__(name=name)

        if not len(bands) == len(offsets) == len(standards):
            raise ValueError("Lengths of bands, standards, and offsets " "must match.")

        self._bands = [get_bandpass(band) for band in bands]
        self._standards = [get_magsystem(s) for s in standards]
        self._offsets = offsets

    @property
    def bands(self):
        return self._bands

    @property
    def standards(self):
        return self._standards

    @property
    def offsets(self):
        return self._offsets

    def _refspectrum_bandflux(self, band):
        if band not in self._bands:
            raise ValueError("band not in local magnitude system")
        i = self._bands.index(band)
        standard = self._standards[i]
        offset = self._offsets[i]

        return 10.0 ** (0.4 * offset) * standard.zpbandflux(band)

    def __str__(self):
        s = "CompositeMagSystem {!r}:\n".format(self.name)

        for i in range(len(self._bands)):
            s += "  {!r}: system={!r}  offset={}\n".format(
                self._bands[i].name, self._standards[i].name, self._offsets[i]
            )
        return s

class SpectralMagSystem(MagSystem):
    """A magnitude system defined by a fundamental spectrophotometric
    standard.

    Parameters
    ----------
    refspectrum : `sncosmo.Spectrum`
        The spectrum of the fundamental spectrophotometric standard.
    """

    def __init__(self, refspectrum, name=None):
        super(SpectralMagSystem, self).__init__(name)
        self._refspectrum = refspectrum

    def _refspectrum_bandflux(self, band):
        return self._refspectrum.bandflux(band)

class ABMagSystem(MagSystem):
    """Magnitude system where a source with F_nu = 3631 Jansky at all
    frequencies has magnitude 0 in all bands."""

    def _refspectrum_bandflux(self, band):
        bwave, btrans = band.to_unit(u.Hz)

        # AB spectrum is 3631 x 10^{-23} erg/s/cm^2/Hz
        # Get spectral values in photons/cm^2/s/Hz at bandpass wavelengths
        # by dividing by (h \nu).
        f = 3631.0e-23 / const.h.cgs.value / bwave

        binw = np.gradient(bwave)
        return np.sum(f * btrans * binw)

import numpy,import astropy,from astropy