import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import makeqx as mkq


def plot_qx(q_x, with_dig=True, with_swatch=True):
    fig = plt.subplots(2,1, figsize=(8,2))
    uni = np.linspace(0, q_x.size, q_x.size, endpoint=False)
    q_xd = mkq.digitize_qx(q_x)
    ax = plt.subplot(211)
    if q_x.size >=100:
        gratim = np.tile(q_x, (int(q_x.size/40),1))
    else:
        gratim = np.tile(q_x, (int(q_x.size/4),1))
    #ax.axis('tight')
    ax.axis('off')
    ax.set_axis_off()
    ax.set_xlim(0,q_x.size)
    ax.xaxis.set_label_position('top') 
    ax.xaxis.tick_top()
    ax.imshow(gratim, cmap='binary_r')

    ax1 = plt.subplot(212)
    #ax1.imshow(gratim, cmap='binary_r')
    ax1.plot(uni, q_x, '--', linewidth=1 )
    ax1.plot(uni, q_xd, linewidth=0.8 )
    ax1.set_ylim(-0.1,1.1)
    ax1.set_xlim(0,q_x.size)
    ax1.set_xlabel('Layer Index')
    #ax1.xaxis.set_label_position('top') 
    #ax1.xaxis.tick_top()
    plt.subplots_adjust(wspace=0, hspace=0)
    plt.tight_layout()
    
def plotrsrp(Rs, Rp, lam_low=400, lam_high=1200, lam_pts=100, ang_low=0,
        ang_high=90, ang_pts=25):
    """TODO: Docstring for plotrsrp.

    :Rs: TODO
    :Rp: TODO
    :lam_low: TODO
    :lam_high: TODO
    :lam_pts: TODO
    :ang_low: TODO
    :ang_high: TODO
    :ang_pts: TODO
    :returns: TODO

    """
    fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(10,4))
    for idx, ax in enumerate(axes.flat):
         if idx == 0:  
                im = ax.imshow(Rs, interpolation='bicubic', extent=[lam_low,
                    lam_high, ang_low, ang_high], 
                    vmin=0, vmax=1, cmap="hot", aspect=4)
                ax.set_xlabel('Wavelength (nm)')
                ax.set_ylabel('Angle (Degrees)')
                ax.set_title('s - pol')
         else:
             im = ax.imshow(Rp, interpolation='bicubic', extent=[lam_low, 
                 lam_high, ang_low, ang_high], 
                 vmin=0, vmax=1, cmap="hot", aspect=4)
             ax.set_xlabel('Wavelength (nm)')
             ax.set_ylabel('Angle (Degrees)')
             ax.set_title('p - pol')
    plt.tight_layout()  
    fig.colorbar(im, ax=axes.ravel().tolist(), shrink=0.3)