############################################################
#                                                          #
#       Virtual Laboratory of Geomathematics in Python     #
#                                                          #
#             Calculating the volume of a hill             # 
#                      (30.09.2017)                        #
#                                                          #                
#         Complutense University of Madrid, Spain          #
#                                                          #
#   THIS SCRIPT IS PROVIDED BY THE AUTHORS "AS IS" AND     #
#   CAN BE USED BY ANYONE FOR THE PURPOSES OF EDUCATION    #
#   AND RESEARCH.                                          #
#                                                          #
############################################################
import sys

from pylab import *

import xlrd
import easygui as eg
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

import matplotlib.pyplot as plt
from matplotlib import animation

# Ventana informacion
eg.msgbox("In this script:\n 1) We selected a file with a height profile\n 2) The height profile is plotted\n 3) The solid obtained by rotation of the height profile is also plotted",title='Problem 4', ok_button='Continue')

# Eleccion de fichero
seleccion = ["*.xls","*.xlsx","*.ods"]
fichero=eg.fileopenbox(msg="Please choose a file",title="Data file",default="*",filetypes=seleccion)
print(fichero)

# Funcion de apertura y lectura de la columna n de un fichero excel
def read_excel(fichero,n):
        book = xlrd.open_workbook(fichero)
        sh = book.sheet_by_index(0)

        # Metemos la columna n en una lista
        coln=[]
        for rx in range(1,sh.nrows):
               coln.append(float(sh.cell_value(colx=n-1,rowx=rx)))
        return coln

# Si no se ha seleccionado el fichero se cierra el programa
if fichero == None:
     print("No file selected")
     sys.exit()

global posx,altura
# Lee las posiciones en x
numx = eg.integerbox(msg='Column number for the position (1-100):',
                    title='Position',
                    default=4,
                    lowerbound=1,
                    upperbound=100,
                    image=None)
if numx == None:
       print("No column selected")
       sys.exit()

posx=read_excel(fichero,numx)
posxmax=max(posx)

for i in range(len(posx)):
       posx[i]=posxmax-posx[i]

# Lee las posiciones en y
numy = eg.integerbox(msg='Column number for the height (1-100):',
                    title='Height',
                    default=3,
                    lowerbound=1,
                    upperbound=100,
                    image=None)
if numy == None:
       print("No column selected")
       sys.exit()
altura=read_excel(fichero,numy)


def alt3d(anmax):

   # Precision en los angulos
   n_angles=anmax
   n_radii=len(posx)
   angles=np.linspace(0,anmax*np.pi/180, n_angles, endpoint=True)
   angles= np.repeat(angles[...,np.newaxis], n_radii,axis=1)
   altura3d=[]
   x=(posx*np.cos(angles)).flatten()
   y=(posx*np.sin(angles)).flatten()
   z=np.tile(altura,n_angles)
   z=z.flatten()
   return x,y,z

def alt3dinter(anmax):

   # Precision en los angulos
   n_angles=3
   n_radii=len(posx)
   angles=np.linspace((anmax-5)*np.pi/180,anmax*np.pi/180, n_angles, endpoint=True)
   angles= np.repeat(angles[...,np.newaxis], n_radii,axis=1)
   altura3d=[]
   x=(posx*np.cos(angles)).flatten()
   y=(posx*np.sin(angles)).flatten()
   z=np.tile(altura,n_angles)
   z=z.flatten()
   return x,y,z

# Limites en el eje x de la grafica
minx=int(min(posx))
maxx=int(max(posx))

global altmi,altma
# Limites en altura
altmi=int(min(altura)-50)
altma=int(max(altura)+50)

# Superficie
x,y,altura3d = alt3d(360)

# GRAFICAS
fig1 = plt.figure('Profiles',figsize=plt.figaspect(2.))
fig2 = plt.figure('Animation')

# Grafica de la altura
ax1=fig1.add_subplot(2,1,1)
ax1.tick_params(direction='in')                   # Ticks en la parte interior de la grafica
ax1.yaxis.set_ticks_position('both')              # Ticks en ambos ejes verticales
ax1.xaxis.set_ticks_position('both')              # Ticks en ambos ejes horizontales
ax1.plot(posx,altura,color="blue")                # Dibuja la altura en azul
ax1.set_title(fichero)                            # Nombre del fichero como titulo
ax1.set_ylabel("Height (m)")                      # Etiqueta vertical
ax1.set_xlabel("Radius (m)")                      # Etiqueta vertical
ax1.set_xlim(minx,maxx)                           # Limites de la grafica
ax1.set_ylim(altmi,altma)

ax2=fig1.add_subplot(2,1,2,projection='3d')
ax2.tick_params(direction='in')                   # Ticks en la parte interior de la grafica
ax2.yaxis.set_ticks_position('both')              # Ticks en ambos ejes verticales
ax2.xaxis.set_ticks_position('both')              # Ticks en ambos ejes horizontales
ax2.set_zlabel("Height (m)")
ax2.set_xlabel("X (m)")
ax2.set_ylabel("Y (m)")
ax2.plot([0,0],[0,0],[altmi,altma],'-.',color='k')
ax2.set_zlim(altmi,altma)
ax2.plot_trisurf(x,y,altura3d,cmap=cm.coolwarm,linewidth=0.5)

# Grafica de la animacion
ax3=fig2.add_subplot(1,1,1,projection='3d')
ax3.set_zlim(altmi,altma)
ax3.set_xlim(-maxx,maxx)
ax3.set_ylim(-maxx,maxx)
ax3.set_xlabel('X (m)')
ax3.set_ylabel('Y (m)')
ax3.set_zlabel('Height (m)')

#surf,=ax3.plot([],[],[])

def init():
    surf = ax3.plot([0,0],[0,0],[altmi,altma],'-.',color='k')
    return surf,

# funcion de animacion, llamada secuencialmente
def animate(frame):
    xdata,ydata,zdata = alt3dinter(frame)
    surf = ax3.plot_trisurf(xdata,ydata,zdata,cmap=cm.coolwarm,linewidth=0.5)
    return surf,

nframes = np.arange(0,360,5)

# animador: 
anim = animation.FuncAnimation(fig2,animate,init_func=init,frames=nframes,repeat=False)

plt.show()