#Copyright (c) 2012-2013, David Nero
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#1. Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#2. Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
#ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from __future__ import division

import sys
import datetime
from pylab import *
import scipy.optimize as so
from matplotlib.patches import Ellipse

N = 20 #grid size
L = 4200
n0 = .5 #.5173


def load_train(sky):
    global x, y, e1, e2, N_halos, x_com, y_com, hx, hy
    x,y,e1,e2 = loadtxt('Train_Skies/Training_Sky%i.csv' % sky, skiprows=1,
        delimiter=',', usecols=(1,2,3,4), unpack=True)
    solutions = loadtxt('Training_halos.csv', skiprows=1, delimiter=',', usecols=(1,2,3,4,5,6,7,8,9))
    N_halos, x_com, y_com, hx1, hy1, hx2, hy2, hx3, hy3 = solutions[sky-1].transpose()
    N_halos = int(N_halos)
    hx = r_[hx1, hx2, hx3]
    hy = r_[hy1, hy2, hy3]
    hx = hx[hx.nonzero()]
    hy = hy[hy.nonzero()]


def load_test(sky):
    global x, y, e1, e2, N_halos
    x,y,e1,e2 = loadtxt('Test_Skies/Test_Sky%i.csv' % sky, skiprows=1,
        delimiter=',', usecols=(1,2,3,4), unpack=True)
    N_halos = [1,2,3][(sky-1)//40]

    
def make_grid(e1, e2):
    grid = zeros((N, N), float)
    for yi in range(N):
        yy = yi*L/N
        for xi in range(N):
            xx = xi*L/N
            grid[yi,xi] = et_avg(xx, yy, e1, e2)
    return grid

    
def grid_max(grid):
    return (r_[nonzero(grid==grid.max())]*L/N)[::-1]

    
def fit(x0, y0, r0):
    e1, e2 = halo(x0, y0, r0)
    return make_grid(e1, e2)


def halo(x0, y0, r0):
    if r0 == 0:
        return 0, 0
    else:
        r = sqrt((x-x0)**2 + (y-y0)**2)
        gt = exp(-(r/r0)**n0)
        t = arctan2(y-y0, x-x0)
        e1 = -gt*cos(2*t)
        e2 = -gt*sin(2*t)
        return e1, e2


def et_avg(xx, yy, e1, e2):
    p = arctan2(y-yy, x-xx)
    et = -e1*cos(2*p) - e2*sin(2*p)
    return mean(et)


def gt(x0, y0, n):
    e1_1, e2_1 = halo(x1, y1, r1)
    e1_2, e2_2 = halo(x2, y2, r2)
    e1_3, e2_3 = halo(x3, y3, r3)
    e1_new = e1 - [e1_2+e1_3, e1_1+e1_3, e1_1+e1_2][n-1]
    e2_new = e2 - [e2_2+e2_3, e2_1+e2_3, e2_1+e2_2][n-1]
    return et_avg(x0, y0, e1_new, e2_new)

    
def locate_halo(params, n):
    x0, y0 = params
    if 0 < x0 < L and 0 < y0 < L:
        return 1 - gt(x0, y0, n)
    else:
        return 1e30

def new_fitter(r0, x0, y0, n):
    r = sqrt((x-x0)**2 + (y-y0)**2)
    gt_data = gt(x0, y0, n)
    gt_fit = mean(exp(-(r/r0)**n0))
    return gt_data - gt_fit


#def gt_measured(x0, y0):
    #t = arctan2(y-y0, x-x0)
    #gt = -e1*cos(2*t) - e2*sin(2*t)
    #r = sqrt((x-x0)**2 + (y-y0)**2)
    #a = rec.fromarrays([r, gt], names='r,gt')
    #a.sort(order=['r'])
    #r = a.r
    #gt = a.gt
    #return r, gt, a

    
def find_halo(n):
    diff = grid - [fitted2+fitted3, fitted1+fitted3, fitted1+fitted2][n-1]
    guess = grid_max(diff)
    #guess = (hx[n-1], hy[n-1])
    x0, y0 = so.fmin(locate_halo, guess, args=(n,), disp=False)
    return x0, y0, diff

    
#def callback(params):
    #print params

    
def fit_halo(n):
    x0, y0, diff = find_halo(n)
    guess = [(r1,), (r2,), (r3,)][n-1]
    #r0 = so.fmin(new_fitter, guess, args=(x0,y0,n), disp=False)
    r0 = so.brentq(new_fitter, 0, 1000, args=(x0,y0,n))
    fitted = fit(x0, y0, r0)
    print '        Halo %i: (%4i, %4i), scale: %4i' % (n,x0,y0,r0)
    return x0, y0, r0, fitted, diff


def plot_all():
    N_plots = 5
    fig = range(N_plots)
    ax = range(N_plots)
    im = range(N_plots)
    for i in range(N_plots):
        title = ["Tangential Ellipticity",
                "Halo 1",
                "Halo 2",
                "Halo 3",
                "Residual"]
        fn = [grid, diff1, diff2, diff3, delta]
        limits = [(None,None), (None,None), (None,None), (None,None), (-grid.max(),grid.max())]

        if fn[i].max() != fn[i].min():
            fig[i] = figure(i)
            ax[i] = fig[i].add_subplot(111, aspect='equal')
            ax[i].set_title(title[i])
            vmin, vmax = limits[i]
            im[i] = ax[i].imshow(fn[i], origin='lower', extent=(0,L,0,L), interpolation='nearest', vmin=vmin, vmax=vmax)
            fig[i].colorbar(im[i])
            if type_ != 'test':
                ax[i].scatter(hx, hy, s=100, marker='o', facecolor='black', edgecolor='white')
            ax[i].scatter(fit_hx[:N_halos], fit_hy[:N_halos], s=100, marker='d', facecolor='black', edgecolor='white')
            ax[i].set_xlim(0,L)
            ax[i].set_ylim(0,L)

    try:
        ax[0].contour(bestfit, extent=(0,L,0,L), colors='black')
    except ValueError:
        print 'Unable to draw contours'
    if N_halos >= 1:
        try:
            ax[1].contour(fitted1, extent=(0,L,0,L), colors='black')
        except ValueError:
            print 'Unable to draw contours'
    if N_halos >= 2:
        try:
            ax[2].contour(fitted2, extent=(0,L,0,L), colors='black')
        except ValueError:
            print 'Unable to draw contours'
    if N_halos >= 3:
        try:
            ax[3].contour(fitted3, extent=(0,L,0,L), colors='black')
        except ValueError:
            print 'Unable to draw contours'

    e = sqrt(e1**2+e2**2)
    th = .5*arctan2(e2, e1)
    a = 1/(1-e)
    b = 1/(1+e)

    galaxies = [Ellipse(xy=(x[i],y[i]), width=50*a[i]*sqrt(1-e[i]**2),
        height=50*b[i]*sqrt(1-e[i]**2), angle=th[i]*180/pi, fc='white', ec='black', alpha=.5)
        for i in range(len(x))]

    for g in galaxies:
        ax[0].add_artist(g)

    show()


def plot_less(sky):
    show()
    fig = figure()
    fig.clear()
    ax = fig.add_subplot(111, aspect='equal')
    ax.set_title('Sky %i - Score: %f' % (sky, score))
    im = ax.imshow(grid, origin='lower', extent=(0,L,0,L), interpolation='nearest')
    fig.colorbar(im)
    if type_ != 'test':
        ax.scatter(hx, hy, s=100, marker='o', facecolor='black', edgecolor='white')
    ax.scatter(fit_hx[:N_halos], fit_hy[:N_halos], s=100, marker='d', facecolor='black', edgecolor='white')
    try:
        ax.contour(bestfit, extent=(0,L,0,L), colors='black')
    except ValueError:
        print 'Unable to draw contours'
    ax.set_xlim(0,L)
    ax.set_ylim(0,L)
        
    e = sqrt(e1**2+e2**2)
    th = .5*arctan2(e2, e1)
    a = 1/(1-e)
    b = 1/(1+e)

    galaxies = [Ellipse(xy=(x[i],y[i]), width=50*a[i]*sqrt(1-e[i]**2),
        height=50*b[i]*sqrt(1-e[i]**2), angle=th[i]*180/pi, fc='white', ec='black', alpha=.5)
        for i in range(len(x))]

    for g in galaxies:
        ax.add_artist(g)

    draw()
    

def iterate():
    global grid, fitted1, fitted2, fitted3, diff1, diff2, diff3, bestfit, delta
    global x1, y1, r1, x2, y2, r2, x3, y3, r3
    global fit_hx, fit_hy, score
    grid = make_grid(e1, e2)
    fitted1 = zeros((N, N), float)
    fitted2 = zeros((N, N), float)
    fitted3 = zeros((N, N), float)
    diff1 = zeros((N, N), float)
    diff2 = zeros((N, N), float)
    diff3 = zeros((N, N), float)

    x1 = x2 = x3 = 0
    y1 = y2 = y3 = 0
    r1 = r2 = r3 = 0

    iterations = [1,2,3][N_halos-1]
    #iterations = 2
    for i in range(iterations):
        print 'Iteration %i of %i:' % (i+1, iterations)
        r1 = 300
        x1, y1, r1, fitted1, diff1 = fit_halo(1)
        
        if N_halos >= 2:
            r2 = 300
            x2, y2, r2, fitted2, diff2 = fit_halo(2)
        if N_halos >= 3:
            r3 = 300
            x3, y3, r3, fitted3, diff3 = fit_halo(3)
        fit_hx = r_[x1, x2, x3]
        fit_hy = r_[y1, y2, y3]

        bestfit = fitted1 + fitted2 + fitted3
        delta = grid - bestfit
    score = delta.std()/grid.max()
    print '        Max et: %f Score: %f' % (grid.max(), score)
    

if __name__ == '__main__':
    type_ = sys.argv[1]
        
    if type_ == 'train':
        sky = sys.argv[2]
        if sky != 'all':
            sky = int(sky)
            load_train(sky)
            iterate()
            plot_all()
        else:
            out= ['SkyId,halo_x1,halo_y1,halo_x2,halo_y2,halo_x3,halo_y3\n']
            for sky in range(1,301):
            #for sky in [1,2,3,4,5,20,39,69,75,82]:
                print '*************** Sky %i ***************' % sky
                load_train(sky)
                iterate()
                #if grid.max() < .03:
                    #plot_less(sky)
                out.append('Sky%i,%.1f,%.1f,%.1f,%.1f,%.1f,%.1f\n' % (sky, x1, y1, x2, y2, x3, y3))

            with open('my_training_halos.csv', 'w') as outfile:
                outfile.writelines(out)

    elif type_ == 'test':
        sky = sys.argv[2]
        if sky != 'all':
            sky = int(sky)
            load_test(sky)
            iterate()
            plot_all()

        else:
            out= ['SkyId,halo_x1,halo_y1,halo_x2,halo_y2,halo_x3,halo_y3\n']
            for sky in range(1,121):
                print '*************** Sky %i ***************' % sky
                load_test(sky)
                iterate()
                out.append('Sky%i,%.1f,%.1f,%.1f,%.1f,%.1f,%.1f\n' % (sky, x1, y1, x2, y2, x3, y3))

            today = datetime.datetime.today()
            year = today.year
            month = today.month
            day = today.day
            hour = today.hour
            minute = today.minute
            filename = 'halos_%04i_%02i_%02i_%02i%02i.csv' % (year, month, day, hour, minute)
            with open(filename, 'w') as outfile:
                outfile.writelines(out)

    elif type_ == 'debug':
        skies = [int(i) for i in sys.argv[2:]]
        for n0 in arange(5160,5180)*.0001:
            dr = []
            for sky in range(1,101):
                load_train(sky)
                x0 = hx[0]
                y0 = hy[0]
                r, gt, a = gt_measured(x0, y0)
                #gt_smooth = convolve(gt, ones(100)/100, 'same')
                F = lambda r, r0: exp(-(r/r0)**n0)
                r0, dr0 = so.curve_fit(F, r, gt, (100))
                dr.append(sqrt(dr0)/r0)
                #print sky, r0, err
                #figure()
                #title('Sky %i' % sky)
                #plot(r, gt, 'g.')
                #plot(r, gt_smooth, 'b-')
                #plot(r, F(r,r0), 'k-')
                ##errorbar(rr, ggt, yerr=err, fmt='b.-')
                #show()
            print 'n0 = %f' % n0, mean(dr)