from visual import *

class Surface:
    def __init__(self,domain,scale=1.0/1.0):
        """Create visualisation of domain
        """

        self.frame  = frame()
        self.bed_model     = faces(frame=self.frame)
        self.stage_model = faces(frame=self.frame)
        self.domain = domain
        self.scale  = scale

        self.vertices = domain.vertex_coordinates
        self.bed   = domain.quantities['elevation'].vertex_values
        self.stage = domain.quantities['stage'].vertex_values

        self.max_x = max(max(self.vertices[:,0],self.vertices[:,2],self.vertices[:,4]))
        self.min_x = min(min(self.vertices[:,0],self.vertices[:,2],self.vertices[:,4]))
        self.max_y = max(max(self.vertices[:,1],self.vertices[:,3],self.vertices[:,5]))
        self.min_y = min(min(self.vertices[:,1],self.vertices[:,3],self.vertices[:,5]))
        self.range_x = self.max_x - self.min_x
        self.range_y = self.max_y - self.min_y

        self.range_xy = max(self.range_x, self.range_y)

        self.max_bed = max(self.bed)
        self.min_bed = min(self.bed)
        self.range_bed = self.max_bed - self.min_bed

        print self.max_x, self.min_x, self.max_y, self.min_y, self.min_bed
        print self.max_bed, self.range_x, self.range_y, self.range_xy, self.range_bed

        print 'Calculating triangles'

        self.pos     = zeros( (6*len(domain),3), Float)
        self.colour  = zeros( (6*len(domain),3), Float)
        self.normals = zeros( (6*len(domain),3), Float)

        print 'keys',self.domain.quantities.keys()
        print 'shape of stage',shape(self.stage)

        self.update_all()

    def update_all(self):

        self.update_bed()
        self.update_stage()

    def update_bed(self):

        print 'update bed arrays'
        c0 = 0.3
        c1 = 0.3
        c2 = 0.3
        self.update_arrays(self.bed,  (c0,c1,c2) )
        print 'update bed image'
        self.bed_model.pos    = self.pos
        self.bed_model.color  = self.colour
        self.bed_model.normal = self.normals


    def update_stage(self):
        c0 = 0.1
        c1 = 0.4
        c2 = 0.99
        print 'update stage arrays'
        self.update_arrays(self.stage, (c0,c1,c2) )
        print 'update stage image'
        self.stage_model.pos    = self.pos
        self.stage_model.color  = self.colour
        self.stage_model.normal = self.normals

    def calc_normal(self,v):

        from math import sqrt

        cross = zeros(3,Float)
        normal = zeros(3,Float)

        v10 = v[1,:]-v[0,:]
        v20 = v[2,:]-v[0,:]

        normal[0] = v10[1]*v20[2] - v20[1]*v10[2]
        normal[1] = v10[2]*v20[0] - v20[2]*v10[0]
        normal[2] = v10[0]*v20[1] - v20[0]*v10[1]

        norm = sqrt( normal[0]**2 + normal[1]**2 + normal[2]**2)

        normal[0] = normal[0]/norm
        normal[1] = normal[1]/norm
        normal[2] = normal[2]/norm

        return normal



    def update_arrays(self, quantity, qcolor):

        col = asarray(qcolor)
        v = zeros( (3,3), Float)
        normal = zeros( 3, Float)
        N = len(self.domain)
        scale = self.scale



        for i in range( N ):

            for j in range(3):
                v[j,0] = 2.0*(self.vertices[i,2*j  ]-self.min_x)/self.range_xy-1.0
                v[j,1] = 2.0*(self.vertices[i,2*j+1]-self.min_y)/self.range_xy-1.0
                v[j,2] = quantity[i,j]*scale



            normal = self.calc_normal(v)

            self.pos[6*i  ,:] = v[0,:]
            self.pos[6*i+1,:] = v[1,:]
            self.pos[6*i+2,:] = v[2,:]
            self.pos[6*i+3,:] = v[0,:]
            self.pos[6*i+4,:] = v[2,:]
            self.pos[6*i+5,:] = v[1,:]



            self.colour[6*i  ,:] = col
            self.colour[6*i+1,:] = col
            self.colour[6*i+2,:] = col
            self.colour[6*i+3,:] = col
            self.colour[6*i+4,:] = col
            self.colour[6*i+5,:] = col

            self.normals[6*i  ,:] = normal
            self.normals[6*i+1,:] = normal
            self.normals[6*i+2,:] = normal
            self.normals[6*i+3,:] = -normal
            self.normals[6*i+4,:] = -normal
            self.normals[6*i+5,:] = -normal




#scene.width = 1000
#scene.height = 800

#Original
#scene.center = (0.5,0.5,0)
#scene.forward = vector(-0.1, 0.5, -0.5)

#Temporary (for bedslope)
#scene.forward = vector(0.0006, 0.7, -0.03)


#Temporary for hackett - begin
#scene.autoscale = 0
#scene.scale = (0.002, 0.002, 0.01) #Scale z so that countours stand out more
#scene.center = (300.0,500.0,-10)
#Temporary for hackett - end


#scene.ambient = 0.4
#scene.lights = [(0.6, 0.3, 0.2), (0.1, -0.5, 0.4), (-0.1, 0.1, -0.4),
#               (-0.2, 0.2, 0.1)]


def create_surface(domain):

    surface = Surface(domain)
    domain.surface = surface

def update(domain):

    surface = domain.surface
    surface.update_stage()