# -*- coding: utf-8 -*-
"""
Created on Mon Dec 18 16:31:36 2017

@author: david
"""

import numpy
import scipy
import matplotlib.pyplot as plt
from matplotlib import cm, colors
from mpl_toolkits.mplot3d import Axes3D
import math
import time


def sph2cart(r, theta, phi):
    x = r * numpy.cos(theta) * numpy.sin(phi)
    y = r * numpy.sin(theta) * numpy.sin(phi)
    z = r * numpy.cos(phi)
    
    return x, y, z

def cart2sph(x, y, z):
    r = numpy.sqrt(x**2+y**2+z**2)
    theta = numpy.arctan2(y,x)
    phi = numpy.arccos(z/r)
    #if(x == 0):
    #    phi = 0
    #else:
    #    phi = numpy.arccos(z/r)
    #phi = numpy.arctan2(numpy.sqrt(x**2 + y**2), z)
    return r, theta, phi


def sh(theta, phi, l, m):
    
    if m < 0:
        return numpy.sqrt(2) * (-1)**m * scipy.special.sph_harm(abs(m), l, theta, phi).imag
    elif m > 0:
        return numpy.sqrt(2) * (-1)**m * scipy.special.sph_harm(m, l, theta, phi).real
    else:
        return scipy.special.sph_harm(0, l, theta, phi).real
    
#calculate a spherical harmonic value from a set of coefficients and coordinates P = (theta, phi)    
def sh_coeff(P, C):
    
    s = numpy.zeros(P[0].shape)
    c = range(len(C))
    
    for c in range(len(C)):
        l, m = i2lm(c)                  #get the 2D indices
        s = s + C[c] * sh(P[0], P[1], l, m)
        
    return s

#plot a spherical harmonic function on a sphere using N points
def sh_plot(C, N):       
    phi = numpy.linspace(0, numpy.pi, N)
    theta = numpy.linspace(0, 2*numpy.pi, N)
    phi, theta = numpy.meshgrid(phi, theta)
    
    # The Cartesian coordinates of the unit sphere
    x = numpy.sin(phi) * numpy.cos(theta)
    y = numpy.sin(phi) * numpy.sin(theta)
    z = numpy.cos(phi)
    
    # Calculate the spherical harmonic Y(l,m) and normalize to [0,1]
    fcolors = sh_coeff(theta, phi, C)
    fmax, fmin = fcolors.max(), fcolors.min()
    fcolors = (fcolors - fmin)/(fmax - fmin)
    
    # Set the aspect ratio to 1 so our sphere looks spherical
    fig = plt.figure(figsize=plt.figaspect(1.))
    ax = fig.add_subplot(111, projection='3d')
    ax.plot_surface(x, y, z,  rstride=1, cstride=1, facecolors=cm.seismic(fcolors))
    # Turn off the axis planes
    ax.set_axis_off()
    plt.show()
    
def i2lm(i):
    l = numpy.floor(numpy.sqrt(i))
    m = i - l *(l + 1)
    return l, m

def lm2i(l, m):
    return l * (l+1) + m

#generates a set of spherical harmonic coefficients from samples using linear least squares
def linfit(P, s, nc, clock=False):
    if clock:
        start_time = time.time()
        
    #allocate space for the matrix and RHS values
    A = numpy.zeros((nc, nc))
    b = numpy.zeros(nc)
    
    #calculate each of the matrix coefficients
        #(see SH technical report in the vascular_viz repository)
    for i in range(nc):
        li, mi = i2lm(i)
        yi = sh(P[0], P[1], li, mi)
        for j in range(nc):        
            lj, mj = i2lm(j)
            yj = sh(P[0], P[1], lj, mj)
            A[i, j] = numpy.sum(yi * yj)
        b[i] = numpy.sum(yi * s)            #calculate the RHS value
    
    #calculate the RHS values
    #for j in range(nc):
    #    lj, mj = i2lm(j)
    #    yj = sh(theta, phi, lj, mj)
    #    b[j] = numpy.sum(yj * s)
    
    if clock:
        print("SH::linfit:matrix "+str(time.time() - start_time)+"s")
    #solve the system of linear equations
    R = numpy.linalg.solve(A, b)
    
    if clock:
        print("SH::linfit:solution "+str(time.time() - start_time)+"s")
    return R

#generate a scatter plot in 3D using spherical coordinates
def scatterplot3d(P):
    r, theta, phi = P
    #convert all of the samples to cartesian coordinates
    X, Y, Z = sph2cart(r, theta, phi)
    
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(X, Y, Z)
    plt.show()