# Copyright (C) 2013 Imperial College London
# Written by David Ham

import dolfin
import numpy as np
from numpy import sin,cos

class TorusMesh(dolfin.Mesh):
    """Structured triangular mesh of a torus."""
    def __init__(self, major_resolution, minor_resolution, R=1.0, r=0.25):
        
        dolfin.Mesh.__init__(self)
        self.major_radius = R
        self.minor_radius = r

        e = dolfin.MeshEditor()
        
        e.open(self, 2, 3)
        e.init_vertices(minor_resolution*major_resolution)
        e.init_cells(2*minor_resolution*major_resolution)


        Mr = major_resolution
        mr = minor_resolution
        for i,theta in enumerate(np.linspace(0, 2*np.pi, Mr, endpoint=False)):
            for j,phi in enumerate(np.linspace(0, 2*np.pi, mr, endpoint=False)):
                e.add_vertex(i*mr + j, (R+r*cos(phi))*cos(theta), (R+r*cos(phi))*sin(theta), r*sin(phi))
                
                e.add_cell(2*(i*mr + j), np.array(sorted([i*mr + j, 
                                                          i*mr + (j+1) % mr, 
                                                          ((i+1) % Mr)*mr + j]), 
                                                  dtype=np.uintp))
                e.add_cell(2*(i*mr + j)+1, np.array(sorted([i*mr + (j+1) % mr, 
                                                            ((i+1) % Mr)*mr + j, 
                                                            ((i+1) % Mr)*mr + (j+1) % mr]), 
                                                    dtype=np.uintp))
        e.close()

        self.global_normal = self.polar_expression(dolfin.Expression(("0", "0", "1")))
        
        self.init_cell_orientations(self.global_normal)
        


    def polar_expression(self, in_expression):
        """Given an expression which expects positions in polar
        (theta,phi,r) coordinates, return one which expects expressions in
        Cartesian coordinates. This currently uses the Python
        Expression subclass route, so there is no promise that this is
        particularly fast."""

        if len(in_expression.value_shape()) not in [0,1,3]:
            print in_expression.value_shape()
            raise ValueError("Only know how to deal with value_size [1] and [3] expressions")

        R = self.major_radius
        r = self.minor_radius

        class PolarExpression(dolfin.Expression):
            def eval(self, value, x):
                rhat = (x[0]**2+x[1]**2)**.5 
                theta = np.arctan2(x[1], x[0])
                phi = np.arctan2(x[2], rhat-R)

                if len(value) in (0,1):
                    in_expression.eval(value, np.array([theta, phi, r]))
                else:
                    in_expression.eval(value, np.array([theta, phi, r]))

                    T = np.array([[-(r*cos(phi) + R)*sin(theta), -r*sin(phi)*cos(theta), cos(phi)*cos(theta)], 
                                  [(r*cos(phi) + R)*cos(theta),  -r*sin(phi)*sin(theta), sin(theta)*cos(phi)], 
                                  [0, r*cos(phi), sin(phi)]])
                    value[:] = np.dot(T,value)
                    
            def value_shape(self):
                return in_expression.value_shape()
            
        return PolarExpression()

if __name__=="__main__":
    
    t = TorusMesh(20,10)
    
    F = dolfin.VectorFunctionSpace(t, "CG", 1, dim=3)
    R = dolfin.FunctionSpace(t, "RT", 1)
    
    r=dolfin.project(t.polar_expression(dolfin.Expression(("0","1","0"))),R)

    f=dolfin.project(r,F)
    
    out = dolfin.File("test.pvd")
    
    out << f
