"""Compare the H1 observations of the Antennae to the simulation"""

import matplotlib.pyplot as plt
import numpy as np

from analysis import utils

# These paremeters were checked by running an interactive version of this function
def plot2d(i=135, theta=1.86, phi=4.10, view=5.10, scaling=12, v_scaling=150, v_offset=43, x_offset=-4, y_offset=2):
    """Plot an encounter

    Parameters:
        i (int): timestep
        theta (float): polar angle
        phi (float): azimuthal angle
        view (float): rotation along line of sight
        scaling (float): scaling of the position vectors
        v_scaling (float): scaling of the line of sight velocity
        v_offset (float): offset of the line of sight velocity
        x_offset (float): position offset in the x direction
        y_offset (float): position offset in the y direction
    """
    data = utils.loadData('antennae', 13500)

    r_vec, v_vec = data['r_vec'], data['v_vec']*v_scaling
    mask = data['type'][:,0] == 'disk'
    r_vec, v_vec = r_vec[mask], v_vec[mask]
    
    
    # Viewing angles
    M1 = np.array([[1, 0, 0],
        [0, np.cos(theta), np.sin(theta)],
        [0, -np.sin(theta), np.cos(theta)]])
    M2 = np.array([[np.cos(phi), np.sin(phi), 0],
        [-np.sin(phi), np.cos(phi), 0],
        [0, 0, 1]])
    M3 = np.array([[np.cos(view), np.sin(view), 0],
        [-np.sin(view), np.cos(view), 0],
        [0, 0, 1]])
    
    M = np.matmul(M2, np.matmul(M1, M3))
    r_vec = np.tensordot(r_vec, M, axes=[1,0])
    v_vec = np.tensordot(v_vec, M, axes=[1,0])
    
    # Offsets
    r_vec[:,0] = r_vec[:,0]*scaling + x_offset
    r_vec[:,1] = r_vec[:,1]*scaling + y_offset
    v_vec[:,2] += v_offset

    f, axs = plt.subplots(2, 2,  figsize=((11+14/1.97)/2, (14+11/1.54)/2), 
        gridspec_kw = {'width_ratios':[11, 14/1.97], 'height_ratios':[14, 11/1.54]})
    
    # Plot background images
    # We manually derive the region that must be plotted based on
    # the data from Hibbard
    axs[0,0].imshow(plt.imread('literature/figs/hibbard1.png'), 
        extent=[-58.2, 88.6, -86.5, 78.7])
    axs[1,0].imshow(plt.imread('literature/figs/hibbard3.png'), 
        extent=[-58.2, 88.6, -270, 310])
    axs[0,1].imshow(plt.imread('literature/figs/hibbard2.png'), 
        extent=[-270, 310, -86.5, 78.7])
    
    # Plot
    axs[0,0].scatter(r_vec[:,0], r_vec[:,1], s=3, edgecolor='', alpha=3E-1)
    axs[1,0].scatter(r_vec[:,0], v_vec[:,2], s=3, edgecolor='', alpha=3E-1)
    axs[0,1].scatter(v_vec[:,2], r_vec[:,1], s=3, edgecolor='', alpha=3E-1)
    f.subplots_adjust(hspace=0, wspace=0)
    
    # Get the correct scales on all axes
    axs[0,0].set_xlim((-58.2, 88.6))
    axs[0,0].set_ylim((-86.5, 78.7))
    axs[1,0].set_xlim((-58.2, 88.6))
    axs[1,0].set_ylim((-270, 310))
    axs[1,0].set_aspect((88.6 + 58.2)/(270+310) / 1.54)
    axs[0,1].set_ylim((-86.5, 78.7))
    axs[0,1].set_xlim((-270, 310))
    axs[0,1].set_aspect((270+310)/(86.5 + 78.7) * 1.97)
    axs[1,1].axis('off')
    
    # Styling
    utils.stylizePlot(axs.flatten())
    axs[0,1].set_xlabel(r'$v_{los}$ / km s$^{-1}$', fontsize=14)
    axs[1,0].set_ylabel(r'$v_{los}$ / km s$^{-1}$', fontsize=14)
    axs[1,0].set_xlabel(r'$x$ / kpc', fontsize=14)
    axs[0,0].set_ylabel(r'$y$ / kpc', fontsize=14)
    
    plt.show()
    
plot2d()