#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 20 21:18:28 2012

@author: Bernhard Tittelbach, Othmar Gsenger
"""

import numpy as np
import pylab as pl
import scipy.io
import re
import os

w_x = np.random.ranf() * np.pi/2 #rad/s
w_y = np.random.ranf() * np.pi/2 #rad/s
w_z = np.random.ranf() * np.pi/2 #rad/s
filename_export_mat = os.path.join(os.path.curdir,"pd_sim.mat")
filename_export_csv =  os.path.join(os.path.curdir,"pd_sim.csv")
## fun constellations to simulate:
# sun and earth close next to each other and no light from moon
sun = [1,1,1]
earth = [1,8,-5]
moon = [1,1,1]
# sun and earth next to each other, moon at maximum
#~ sun = [1,1,1]
#~ earth = [1,8,-5]
#~ moon = [-1,-1,-1]
# sun hidden behind earth, only moon visible
#~ sun = [1,1,1]
#~ earth = [1,1,1]
#~ moon = [1,-1,0]
#satellite between sun and earth
#~ sun = [1,1,1]
#~ earth = [-1,-1,-1]
#~ moon = [1,-1,0]
#moon and earth opposite to sun
#~ sun = [1,1,1]
#~ earth = [-1,-2,-1]
#~ moon = [-1,2,-1]
sun_intensity = 1.0
#raised earth_reflection_intensity to account not only for light from a single ray but for multiple rays reflected from a surface.
#todo: this should propably be included in the earth model instead, once I find good way to do that
earth_reflection_intensity = 0.50 * sun_intensity # 0 to turn earth off ;-)
moon_reflection_intensity = 0.07 * sun_intensity # 0 to turn moon off ;-)
orbit_height = 310e3 #m
numperiods = 4
sensors = np.matrix([
                    [ 1,0,0],
                    [ -1,0,0],
                    [0,1,0],
                    [0,-1,0],
                    [0,0,1],
                    [0,0,-1]
                    ])

#calculate simduration so we see numperiods periods of the smalest angular speed
w_xyz = np.matrix([w_x,w_y,w_z])
duration_full_period = np.prod(np.pi / np.array(filter(lambda x: x>0.0,w_xyz))) if w_x or w_y or w_z else 1.0
minimum_period = np.min(np.pi / np.array(filter(lambda x: x>0.0,w_xyz))) if w_x or w_y or w_z else 1.0
simduration = numperiods* duration_full_period #in seconds
simtimestep = minimum_period / 10.0 #s
#calculate angle between normalvector of earthdisc and edge of earthdisc as seen from satellite
earth_radius = 6e6 #m
earth_distance =  earth_radius + orbit_height #m
angle_earthdisc_visible = np.arcsin(float(earth_radius) / earth_distance )
assert(angle_earthdisc_visible < np.pi/2.0)
simtimeline = np.arange(0,simduration,simtimestep)
angular_speed_sum_vectors = ( w_xyz.transpose() * simtimeline ).transpose() % (2 *np.pi)
sun = sun / np.linalg.norm(sun)
earth = earth / np.linalg.norm(earth)
moon = moon / np.linalg.norm(moon)
angle_sun_earth =  np.arccos(max(min(np.dot(sun,earth),1.0),-1.0))
angle_sun_moon =  np.arccos(max(min(np.dot(sun,moon),1.0),-1.0))
angle_earth_moon =  np.arccos(max(min(np.dot(earth,moon),1.0),-1.0))
print angle_sun_earth, angle_sun_moon, angle_earth_moon
#print angle_sun_earth /np.pi*180, angle_earthdisc_visible /np.pi*180

print "Simulation Input:"
print "　 Duration : %.4f [s]" % simduration
print "　 Step Size: %.4f [s]" % simtimestep
print "　 Angular Speed x-axis: %.3f [rad/s] period: %.3f [s]" % (w_x, np.pi / w_x)
print "　 Angular Speed y-axis: %.3f [rad/s] period: %.3f [s]" % (w_y, np.pi/ w_y)
print "　 Angular Speed z-axis: %.3f [rad/s] period: %.3f [s]" % (w_z, np.pi/w_z)
print "　 Full period %.3f [s]" % duration_full_period
print "   Earth " + ("Light intensity: %f" % earth_reflection_intensity if earth_reflection_intensity else "is switched OFF")
print "   Moon  " + ("Light intensity: %f" % moon_reflection_intensity if moon_reflection_intensity else "is switched OFF")
print "   Sphere surface, covered by Earth as seen from satellite: %.2f %%" % ((earth_radius**2) / (4.0 * earth_distance**2) * 100.0)
print "   Satellite " + (u"Is (!)" if angle_sun_earth < angle_earthdisc_visible else u"is not") + u" in Earth's shadow (Angle Sun-Earth: %d°)" % (angle_sun_earth /np.pi*180.0)
print "   Moon " + (u"Is (!)" if angle_earth_moon < angle_earthdisc_visible else u"is not") + u" hidden behind the Earth (Angle Earth-Moon: %d°)" % (angle_earth_moon /np.pi*180.0)
print ""


def getRotationMatrix(angles):
    x_angle = angles[0,0]
    y_angle = angles[0,1]
    z_angle = angles[0,2]
    return np.matrix([
      [1,   0,                  0],
      [0,   np.cos(x_angle),   -np.sin(x_angle)],
      [0,   np.sin(x_angle),    np.cos(x_angle)]
    ]) * np.matrix([
      [ np.cos(y_angle),0,np.sin(y_angle)],
      [ 0,              1,0],
      [-np.sin(y_angle),0,np.cos(y_angle)]
    ]) * np.matrix([
      [np.cos(z_angle),-np.sin(z_angle),0],
      [np.sin(z_angle), np.cos(z_angle),0],
      [0,0,1]
    ])

def rotateVector(vec_to_rotate, angles_vector):
    return vec_to_rotate * getRotationMatrix(angles_vector)

def sensor_model_rel_sensitivity_cos(angle):
    if angle < np.pi/2.0:
        return np.cos(angle)
    else:
        return 0.0

sensor_model = sensor_model_rel_sensitivity_cos

def sun_earth_shadow(angle):
    global angle_sun_earth, angle_earthdisc_visible, sensor_model
    if angle_sun_earth < angle_earthdisc_visible:
        return 0.0
    else:
        return sensor_model(angle)

sun_model = sun_earth_shadow

def earthdisc_light_rel_intensity(angle):
    global angle_sun_earth, angle_earthdisc_visible, sensor_model
    if angle_sun_earth < angle_earthdisc_visible:
        return 0.0
    #todo: factor in angle of actually lightened surface depending on angle_sun_earth and the angle between sat orbit and the sun
    elif angle < angle_earthdisc_visible:
        #account for varying reflection of different earth surface materials [0.75 .. 1.25]
        return 1.0 + (np.random.ranf() - 0.5)/4.0
    else:
        return sensor_model(angle - angle_earthdisc_visible)

earth_model = earthdisc_light_rel_intensity

def moon_reflection_from_sun_rel_intensity(angle):
    """ reflected light from moon is a function of angle between sun and moon as seen from satellite if < 90%
        also: moon can be hidden by earthdisc
    """
    global angle_sun_moon, sensor_model, angle_earth_moon
    if angle_earth_moon < angle_earthdisc_visible:
        return 0.0
    elif angle_sun_moon < np.pi/2.0:
        return angle_sun_moon / (np.pi/2.0) * sensor_model(angle)
    else:
        return sensor_model(angle)

moon_model = moon_reflection_from_sun_rel_intensity

def _quote(str):
    return '"'+ re.sub(r'(?!\\)"','\\"',str)  +'"'

def exportCSV(dst_file_path, names, tline, smatrix):
    with open(dst_file_path, "w") as fh:
        fh.truncate()
        fh.write(";".join(map(_quote, ["SimTime"]+names)) +"\n")
        for row in np.vstack((tline,smatrix.T)).T:
            fh.write(";".join(["%f" % x for x in row.T]) +"\n")

senstrace=[]
suntrace=[]
earthtrace=[]
for cur_rot in angular_speed_sum_vectors:
    rot_sensors=rotateVector(sensors,cur_rot)
    angles_to_sun=np.array(np.arccos(np.dot(rot_sensors,np.matrix(sun).T))).flatten()
    angles_to_earth=np.array(np.arccos(np.dot(rot_sensors,np.matrix(earth).T))).flatten()
    angles_to_moon=np.array(np.arccos(np.dot(rot_sensors,np.matrix(moon).T))).flatten()
    light_intensity_per_sensor = \
        sun_intensity * np.array([sun_model(a) for a in angles_to_sun]) + \
        earth_reflection_intensity * np.array([earth_model(a) for a in angles_to_earth]) + \
        moon_reflection_intensity * np.array([moon_model(a) for a in angles_to_moon])
    senstrace.append(np.minimum(1.0, light_intensity_per_sensor))

    suntrace.append(angles_to_sun)
    earthtrace.append(angles_to_earth)

senstrace = np.matrix(senstrace)
suntrace = np.matrix(suntrace)
earthtrace = np.matrix(earthtrace)
legend_names=["Sensor%d %s" % (x+1, sensors[x].flatten()) for x in range(0,6)]

print "Saving to " + filename_export_mat
scipy.io.savemat(filename_export_mat, mdict={"legend":legend_names, "senstrace":senstrace,"simtimeline":simtimeline,"angular_speed_xyz":w_xyz}, oned_as="column")
print "Saving to " + filename_export_csv
exportCSV(filename_export_csv, legend_names, simtimeline, senstrace)


xlabelpos = np.hstack((np.arange(0,simduration, simduration/5),simduration))
xlabelvalues = ["%.3f" % t for t in xlabelpos]
ylabelpos = [0,0.5,1]
ylabelvalues = [u"0%", u"50%", u"100%"]

#~ pl.figure()
#~ pl.plot(simtimeline, senstrace)
#~ pl.xticks(xlabelpos,xlabelvalues,rotation=0)
#~ pl.yticks(ylabelpos, ylabelvalues)
#~ pl.xlabel("[s]")
#~ pl.ylabel("[rad]")
#~ pl.legend(legend_names)

pl.figure()
for f in range(0,3):
  fs = 2*f
  fe = fs + 2
  pl.subplot(3,1,f+1)
  pl.plot(simtimeline, senstrace[:,fs:fe])
  pl.xticks(xlabelpos,xlabelvalues,rotation=0)
  pl.yticks(ylabelpos, ylabelvalues)
  pl.xlabel("[s]")
  pl.ylabel("[Light Intensity]")
  pl.legend(legend_names[fs:fe])

pl.figure()
spncols=2
spnrows=3
for s in range(0,6):
  pl.subplot(spnrows,spncols,s+1)
  pl.plot(simtimeline, senstrace[:,s],"g")
  pl.plot(simtimeline, 1/(2*np.pi)*suntrace[:,s],"y")
  pl.plot(simtimeline, 1/(2*np.pi)*earthtrace[:,s],"b")
  pl.xticks(xlabelpos,xlabelvalues,rotation=0)
  pl.yticks(ylabelpos, ylabelvalues)
  pl.ylabel("[Light Intensity]")
  pl.title(legend_names[s])
  pl.legend(["Sensor Value","Sun Angle", "Earth Angle"])

pl.show()