import numpy as np
import matplotlib.pyplot as plt

import pdb


def phot_plot(annphot_file_list, totphot_file_list,
                  label_list, plot_filename,
                  color_list=None,
                  asym_mag_list=None):
    """
    Make plots of the surface brightness photometry

    Parameters
    ----------
    annphot_file_list : list of strings
        list of file names with annular photometry (generated by surface_phot.py)

    totphot_file_list : list of strings
        list of file names with total enclosed photometry (generated by surface_phot.py)

    label_list : list of strings
        labels associated with the above files (presumably the filter names) that will be used to label each plot

    plot_filename : string
        name of the output file with the plots

    color_list : list of strings (default = None)
        if set, each profile will be plotted with these colors, otherwise they'll be black

    asym_mag_list : list of strings (default = None)
        list of file names with asymptotic magnitudes (generated by surface_phot.py)
        if set, plot a line in the growth curves showing the symptotic magnitudes

    """

    # total files/filters
    n_file = len(annphot_file_list)
    
    # set up the plotting area:
    # need two plots across (profile, growth curve) and len(phot_file_list) down
    fig = plt.figure(figsize=(8, (n_file+1)*3))

    # choose colors
    if color_list is None:
        color_list = ['black'] * n_file

    # save info for plotting everything overlapping at the end
    phot_ann = {j:k for i in label_list for j,k in zip(('r_'+i,'mu_'+i,'mu_err_'+i), (0,0,0))}
    phot_tot = {j:k for i in label_list for j,k in zip(('r_'+i,'mag_'+i,'mag_err_'+i), (0,0,0))}

    # plot for each filter
    for f in range(n_file):

        # ----- read in data
        
        # annulus data
        data = np.loadtxt(annphot_file_list[f])
        # read in column names
        with open(annphot_file_list[f]) as temp:
            lines = temp.readlines()
            cols = lines[0].strip().split(' ')[1:]
        # put data in dictionary
        phot_dict_ann = {cols[i]:data[:,i] for i in range(len(cols))}
        phot_ann['r_'+label_list[f]] = phot_dict_ann['radius']
        phot_ann['mu_'+label_list[f]] = phot_dict_ann['mu']
        phot_ann['mu_err_'+label_list[f]] = phot_dict_ann['mu_err']

        # total data
        data = np.loadtxt(totphot_file_list[f])
        # read in column names
        with open(totphot_file_list[f]) as temp:
            lines = temp.readlines()
            cols = lines[0].strip().split(' ')[1:]
        # put data in dictionary
        phot_dict_tot = {cols[i]:data[:,i] for i in range(len(cols))}
        phot_tot['r_'+label_list[f]] = phot_dict_tot['radius']
        phot_tot['mag_'+label_list[f]] = phot_dict_tot['mag']
        phot_tot['mag_err_'+label_list[f]] = phot_dict_tot['mag_err']
        # asymptotic magnitudes
        if asym_mag_list is not None:
            data = np.loadtxt(asym_mag_list[f])
            with open(asym_mag_list[f]) as temp:
                lines = temp.readlines()
                cols = lines[0].strip().split(' ')[1:]
            temp = {cols[i]:data[i] for i in range(len(cols))}
            phot_tot['asym_mag_'+label_list[f]] = temp['mag']
            phot_tot['asym_mag_err_'+label_list[f]] = temp['mag_err']


        # ----- annulus plot

        # plot profile
        plt.subplot(n_file+1, 2, f*2 + 1)
        plt.errorbar(phot_dict_ann['radius']/60, phot_dict_ann['mu'],
                         yerr=phot_dict_ann['mu_err'],
                         marker='.', color=color_list[f], ms=5, mew=0,
                         linestyle='-',
                         ecolor=color_list[f], capsize=0)

        # axis stuff
        ax = plt.gca()
        ax.set_xlabel('Radius (arcmin)')
        ax.set_ylabel(r'$\mu$ (AB mag/arcsec$^2$)')
        ax.set_ylim((np.nanmax(phot_dict_ann['mu'])+0.5, np.nanmin(phot_dict_ann['mu'])-0.5))

        # label the plot
        plt.text(0.95, 0.9, label_list[f], ha='right', va='center', transform=ax.transAxes)
        
        plt.tight_layout()

        
        # ----- total plot

        # plot profile
        plt.subplot(n_file+1, 2, f*2 + 2)
        plt.errorbar(phot_dict_tot['radius']/60, phot_dict_tot['mag'],
                         yerr=phot_dict_tot['mag_err'],
                         marker='.', color=color_list[f], ms=5, mew=0,
                         linestyle='-',
                         ecolor=color_list[f], capsize=0)

        
        # axis stuff
        ax = plt.gca()
        ax.set_xlabel('Radius (arcmin)')
        ax.set_xlim(left=0)
        ax.set_ylabel('AB mag')
        ax.set_ylim((np.nanmax(phot_dict_tot['mag'])+0.5, np.nanmin(phot_dict_tot['mag'])-0.5))

        # asymptotic magnitudes
        if asym_mag_list is not None:
            plt.plot(ax.get_xlim(),
                         [phot_tot['asym_mag_'+label_list[f]], phot_tot['asym_mag_'+label_list[f]]],
                         marker='.', color=color_list[f], ms=0, mew=0,
                         linestyle=':')

        # label the plot
        #plt.text(0.95, 0.9, label_list[f], ha='right', va='center', transform=ax.transAxes)
        
        plt.tight_layout()



    # now plots for all of the things

    # annulus
    plt.subplot(n_file+1, 2, n_file*2 + 1)
    # - error regions
    for i in range(n_file):
        #plt.fill_between(phot_ann['r_'+label_list[i]]/60,
        #                     phot_ann['mu_'+label_list[i]] + phot_ann['mu_err_'+label_list[i]],
        #                     phot_ann['mu_'+label_list[i]] - phot_ann['mu_err_'+label_list[i]],
        #                     facecolor=color_list[i], edgecolor='none', alpha=0.2)
        plt.errorbar(phot_ann['r_'+label_list[i]]/60 + i*0.01-((n_file-1)/2*0.01),
                         phot_ann['mu_'+label_list[i]],
                         yerr=phot_ann['mu_err_'+label_list[i]],
                         marker='.', color=color_list[i], ms=0, mew=0,
                         linestyle='none',
                         ecolor=color_list[i], capsize=0, alpha=0.2)
    # - photometry
    for i in range(n_file):
        plt.plot(phot_ann['r_'+label_list[i]]/60, phot_ann['mu_'+label_list[i]],
                     marker='.', color=color_list[i], ms=5, mew=0,
                     linestyle='-')
        
        ymin = np.nanmin(phot_ann['mu_'+label_list[i]])
        ymax = np.nanmax(phot_ann['mu_'+label_list[i]])
        if i == 0:
            yrange = [ymin, ymax]
        if i > 0:
            if ymin < yrange[0]:
                yrange[0] = ymin
            if ymax > yrange[1]:
                yrange[1] = ymax
    # - axis stuff
    ax = plt.gca()
    ax.set_xlabel('Radius (arcmin)')
    ax.set_ylabel(r'$\mu$ (AB mag/arcsec$^2$)')
    ax.set_ylim((yrange[1]+0.5, yrange[0]-0.5))
    plt.tight_layout()


    # total
    plt.subplot(n_file+1, 2, n_file*2 + 2)
    # - error regions
    for i in range(n_file):
        #plt.fill_between(phot_tot['r_'+label_list[i]]/60,
        #                     phot_tot['mag_'+label_list[i]] + phot_tot['mag_err_'+label_list[i]],
        #                     phot_tot['mag_'+label_list[i]] - phot_tot['mag_err_'+label_list[i]],
        #                     facecolor=color_list[i], edgecolor='none', alpha=0.2)
        plt.errorbar(phot_tot['r_'+label_list[i]]/60 + i*0.01-((n_file-1)/2*0.01),
                         phot_tot['mag_'+label_list[i]],
                         yerr=phot_tot['mag_err_'+label_list[i]],
                         marker='.', color=color_list[i], ms=0, mew=0,
                         linestyle='none',
                         ecolor=color_list[i], capsize=0, alpha=0.2)
    # - photometry
    for i in range(n_file):
        plt.plot(phot_tot['r_'+label_list[i]]/60, phot_tot['mag_'+label_list[i]],
                     marker='.', color=color_list[i], ms=5, mew=0,
                     linestyle='-')
        
        ymin = np.nanmin(phot_tot['mag_'+label_list[i]])
        ymax = np.nanmax(phot_tot['mag_'+label_list[i]])
        if i == 0:
            yrange = [ymin, ymax]
        if i > 0:
            if ymin < yrange[0]:
                yrange[0] = ymin
            if ymax > yrange[1]:
                yrange[1] = ymax
    # - axis stuff
    ax = plt.gca()
    ax.set_xlabel('Radius (arcmin)')
    ax.set_xlim(left=0)
    ax.set_ylabel('AB mag')
    ax.set_ylim((yrange[1]+0.5, yrange[0]-0.5))
    plt.tight_layout()
    # - asymptotic magnitudes
    if asym_mag_list is not None:
        for i in range(n_file):
            plt.plot(ax.get_xlim(),
                         [phot_tot['asym_mag_'+label_list[i]], phot_tot['asym_mag_'+label_list[i]]],
                         marker='.', color=color_list[i], ms=0, mew=0,
                         linestyle=':')
        

    # save plot
    plt.savefig(plot_filename)
    plt.close()

    #pdb.set_trace()
