import SASModels.sas_models as sas_models
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import sys, lmfit
from mpl_toolkits.axes_grid1 import make_axes_locatable

sas_models.math.n_integration_cuts = 1

I0 = 0.00002
RL = 3000
R0 = 600
phi = 90
SLDspindle = 42.209e-6 #a-Fe2O3 at 1.3414 A
SLDmatrix = 9.459e-6 #H2O at 1.3414 A
sigRL = 0.0

B = 0.1 # T
Ms = 1642 # A/m
V = 4.6e-22
k =  1.3807e-23
T = 300
MsVperKT = Ms*V/k/T

def get_custom_cmap():
    def make_colormap(seq):
        """Return a LinearSegmentedColormap
        seq: a sequence of floats and RGB-tuples. The floats should be increasing
        and in the interval (0,1).
        """
        cdict = {'red': [], 'green': [], 'blue': []}
        for i, item in enumerate(seq):
            pos, r, g, b = item
            cdict['red'].append([pos, r, r])
            cdict['green'].append([pos, g, g])
            cdict['blue'].append([pos, b, b])
        return mcolors.LinearSegmentedColormap('CustomMap', cdict)
    #Nice Coloring:
    c = mcolors.ColorConverter().to_rgb
    custom_colors = [(0, 0, 0, 0),\
             (0.18, 0.05, 0.05, 0.2),\
             (0.28, 0, 0, 1),\
             (0.4, 0.7, 0.85, 0.9),\
             (0.45, 0, 0.75, 0),\
             (0.6, 1, 1, 0),\
             (0.75, 1, 0, 0),\

             (0.92 , 0.6, 0.6, 0.6),\
             (1  , 0.95, 0.95, 0.95)]
    custom_cmap = make_colormap(custom_colors)
    custom_cmap.set_bad(color='black')
    return custom_cmap

if False:
    psi_range = np.linspace(0, np.pi/2, 90)
    prob=[]
    for psi in psi_range:
        prob_val = sas_models.math.langevin_prob(psi, B*MsVperKT, 0.)
        prob.append(prob_val)

    fig, ax = plt.subplots()
    ax.plot(psi_range, prob)
    plt.show()


pix_size = 100
qy = np.linspace(-0.005, 0.005, pix_size)
qz = np.linspace(-0.005, 0.005, pix_size)
Imodel = sas_models.spindle.get_2dimage_langevin(qy, qz, RL, R0, phi,\
                                            SLDspindle, SLDmatrix,\
                                            sigRL, B, MsVperKT)

Imodel = I0*Imodel


fig, ax = plt.subplots()
cmap = get_custom_cmap()
pcm = ax.pcolormesh(qy, qz, Imodel.T,\
                    norm=mcolors.LogNorm(), cmap=cmap)


ax.set_xlabel('$q_y \, / \, \AA^{-1}$')
ax.set_ylabel('$q_z \, / \, \AA^{-1}$')
ax.set_xlim(min(qy), max(qy))
ax.set_ylim(min(qz), max(qz))

divider3 = make_axes_locatable(ax)
cax = divider3.append_axes('right', size="5%", pad=0.05)
cbar = fig.colorbar(pcm, cax=cax)

fig.tight_layout()
fig.savefig('Spindle.png')
plt.show()
