from vtk import *
import time
import sys

# Enable off-screen rendering with MESA
factGraphics = vtkGraphicsFactory()
factGraphics.SetUseMesaClasses(1)
factImage = vtkImagingFactory()
factImage.SetUseMesaClasses(1)

def readvtk(file):
  reader=vtkPolyDataReader()
  reader.SetFileName(file)
  reader.Update()
  return reader.GetOutput()

def getactor(mesh):
  mapper=vtkPolyDataMapper()
  mapper.SetInput(mesh)
  actor=vtkActor()
  actor.SetMapper(mapper)
  return actor

def sshot(renwin,file):
  renwin.Modified()
  renwin.Render()
  winfil=vtkWindowToImageFilter()
  winfil.SetInput(renwin)
  png=vtkPNGWriter()
  png.SetInput(winfil.GetOutput())
  png.SetFileName(file)
  png.Write()


def shift(pdin,dir,amt):
  pd = vtkPolyData();
  pd.DeepCopy(pdin);
  spk = pd.GetPointData().GetNormals();

  for i in range(0,pd.GetNumberOfPoints()):
    px = [0,0,0];

    pd.GetPoint(i, px);
    for j in range(0,2):
      px[j] = px[j] + dir * amt * spk.GetComponent(i,j)
    pd.GetPoints().SetPoint(i, px)

  return pd


def procclust(pd, name):
  dat = pd.GetPointData().GetArray(name);
  pd.GetPointData().SetScalars(dat);
  return pd;

def ctable():
  tab = vtkLookupTable();
  tab.SetNumberOfTableValues(64);
  tab.SetTableValue(0,0,0,0.5625,1);
  tab.SetTableValue(1,0,0,0.6250,1);
  tab.SetTableValue(2,0,0,0.6875,1);
  tab.SetTableValue(3,0,0,0.7500,1);
  tab.SetTableValue(4,0,0,0.8125,1);
  tab.SetTableValue(5,0,0,0.8750,1);
  tab.SetTableValue(6,0,0,0.9375,1);
  tab.SetTableValue(7,0,0,1.0000,1);
  tab.SetTableValue(8,0,0.0625,1.0000,1);
  tab.SetTableValue(9,0,0.1250,1.0000,1);
  tab.SetTableValue(10,0,0.1875,1.0000,1);
  tab.SetTableValue(11,0,0.2500,1.0000,1);
  tab.SetTableValue(12,0,0.3125,1.0000,1);
  tab.SetTableValue(13,0,0.3750,1.0000,1);
  tab.SetTableValue(14,0,0.4375,1.0000,1);
  tab.SetTableValue(15,0,0.5000,1.0000,1);
  tab.SetTableValue(16,0,0.5625,1.0000,1);
  tab.SetTableValue(17,0,0.6250,1.0000,1);
  tab.SetTableValue(18,0,0.6875,1.0000,1);
  tab.SetTableValue(19,0,0.7500,1.0000,1);
  tab.SetTableValue(20,0,0.8125,1.0000,1);
  tab.SetTableValue(21,0,0.8750,1.0000,1);
  tab.SetTableValue(22,0,0.9375,1.0000,1);
  tab.SetTableValue(23,0,1.0000,1.0000,1);
  tab.SetTableValue(24,0.0625,1.0000,0.9375,1);
  tab.SetTableValue(25,0.1250,1.0000,0.8750,1);
  tab.SetTableValue(26,0.1875,1.0000,0.8125,1);
  tab.SetTableValue(27,0.2500,1.0000,0.7500,1);
  tab.SetTableValue(28,0.3125,1.0000,0.6875,1);
  tab.SetTableValue(29,0.3750,1.0000,0.6250,1);
  tab.SetTableValue(30,0.4375,1.0000,0.5625,1);
  tab.SetTableValue(31,0.5000,1.0000,0.5000,1);
  tab.SetTableValue(32,0.5625,1.0000,0.4375,1);
  tab.SetTableValue(33,0.6250,1.0000,0.3750,1);
  tab.SetTableValue(34,0.6875,1.0000,0.3125,1);
  tab.SetTableValue(35,0.7500,1.0000,0.2500,1);
  tab.SetTableValue(36,0.8125,1.0000,0.1875,1);
  tab.SetTableValue(37,0.8750,1.0000,0.1250,1);
  tab.SetTableValue(38,0.9375,1.0000,0.0625,1);
  tab.SetTableValue(39,1.0000,1.0000,0,1);
  tab.SetTableValue(40,1.0000,0.9375,0,1);
  tab.SetTableValue(41,1.0000,0.8750,0,1);
  tab.SetTableValue(42,1.0000,0.8125,0,1);
  tab.SetTableValue(43,1.0000,0.7500,0,1);
  tab.SetTableValue(44,1.0000,0.6875,0,1);
  tab.SetTableValue(45,1.0000,0.6250,0,1);
  tab.SetTableValue(46,1.0000,0.5625,0,1);
  tab.SetTableValue(47,1.0000,0.5000,0,1);
  tab.SetTableValue(48,1.0000,0.4375,0,1);
  tab.SetTableValue(49,1.0000,0.3750,0,1);
  tab.SetTableValue(50,1.0000,0.3125,0,1);
  tab.SetTableValue(51,1.0000,0.2500,0,1);
  tab.SetTableValue(52,1.0000,0.1875,0,1);
  tab.SetTableValue(53,1.0000,0.1250,0,1);
  tab.SetTableValue(54,1.0000,0.0625,0,1);
  tab.SetTableValue(55,1.0000,0,0,1);
  tab.SetTableValue(56,0.9375,0,0,1);
  tab.SetTableValue(57,0.8750,0,0,1);
  tab.SetTableValue(58,0.8125,0,0,1);
  tab.SetTableValue(59,0.7500,0,0,1);
  tab.SetTableValue(60,0.6875,0,0,1);
  tab.SetTableValue(61,0.6250,0,0,1);
  tab.SetTableValue(62,0.5625,0,0,1);
  tab.SetTableValue(63,0.5000,0,0,1);
  tab.SetRampToLinear();
  return tab;

def runme(path_to_stat, path_to_edges, name, sname, rgb, campos, camup, stat, type, threshold):

  # stats
  ttestArrayName="T-statistic";

  # create a rendering window and renderer
  ren = vtkRenderer()
  renWin = vtkRenderWindow()
  renWin.DoubleBufferOn()
  #renWin.OffScreenRenderingOn();
  renWin.AddRenderer(ren)
  renWin.SetSize(1200,800)
  renWin.LineSmoothingOn()
  
  # create an actor and give it cone geometry
  pdMain = readvtk(path_to_stat + "ixi_template_" + name + "_def3.med." + type + ".clusters." + stat +  "_" + "%0.2f"%threshold + ".vtk")
  pdClust = readvtk(path_to_edges + "ixi_template_" + name + "_def3_edges.med." + type + ".clusters." + stat +  "_" + "%0.2f"%threshold + ".vtk")

  # Create wireframes
  am = getactor(pdMain);
  am.GetMapper().ScalarVisibilityOn()
  am.GetMapper().SetScalarModeToUsePointFieldData()
  am.GetMapper().SelectColorArray(ttestArrayName);
  am.GetMapper().SetScalarRange(-3.4,3.4)
  #am.GetMapper().SetLookupTable(ctable());
  am.GetMapper().GetLookupTable().SetHueRange(0.66667,0);

  ren.AddActor(am)


  am2 = getactor(pdClust);
  am2.GetMapper().ScalarVisibilityOn()
  am2.GetMapper().SetScalarModeToUseCellFieldData()
  am2.GetProperty().SetColor(0,0,0);
  am2.GetProperty().SetRepresentationToWireframe();
  am2.GetProperty().SetSpecular(0.0);
  am2.GetProperty().SetAmbient(1.0);
  am2.GetProperty().SetDiffuse(0.0);
  am2.GetProperty().SetLineWidth(2);

  ren.AddActor(am2)

  wf1 = getactor(shift(pdMain,1,0.05))
  wf2 = getactor(shift(pdMain,-1,0.05))
  wf1.GetProperty().SetRepresentationToWireframe();
  wf2.GetProperty().SetRepresentationToWireframe();
  #wf1.GetProperty().SetColor(1,1,0);
  #wf2.GetProperty().SetColor(1,1,0);
  wf1.GetProperty().SetSpecular(0.0);
  wf1.GetProperty().SetDiffuse(0.0);
  wf1.GetProperty().SetAmbient(1.0);
  wf2.GetProperty().SetSpecular(0.0);
  wf2.GetProperty().SetDiffuse(0.0);
  wf2.GetProperty().SetAmbient(1.0);
  wf1.GetProperty().SetLineWidth(.5);
  wf2.GetProperty().SetLineWidth(.5);
  wf1.GetProperty().SetOpacity(.15);
  wf2.GetProperty().SetOpacity(.15);
  ren.AddActor(wf1);
  ren.AddActor(wf2)


  # assign our actor to the renderer
  ren.SetBackground(0.32941,0.34902,0.42745)
  ren.GetActiveCamera().SetPosition(campos[0], campos[1], campos[2])
  ren.GetActiveCamera().SetViewUp(camup[0], camup[1], camup[2])
  ren.ResetCamera()
  ren.GetActiveCamera().Zoom(1.4)
 
  print stat + "_" + type + "_tmap_" + sname + ".png";
  sshot(renWin,stat + "_" + type + "_tmap_" + sname + ".png");

# Do all the individual structures
A=[]
#print sys.argv;
stat=sys.argv[1];
type=sys.argv[2];
threshold=float(sys.argv[3]);
path_to_stats=sys.argv[4]
path_to_edges=sys.argv[5]
#sys.exit();


runme(path_to_stats,path_to_edges, "cc","left_cc",[255,0,  0]  , [1,0,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges, "cc","right_cc",[255,0,  0]  , [-1,0,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges,"cc","top_cc",[255,0,  0]  , [0,0,1], [1,0,0], stat, type, threshold);
runme(path_to_stats,path_to_edges,"cc","bottom_cc",[255,0,  0]  , [0,0,-1], [1,0,0], stat, type, threshold);
runme(path_to_stats,path_to_edges,"cc","front_cc",[255,0,  0]  , [0,1,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges,"cc","back_cc",[255,0,  0]  , [0,-1,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges,"left_unc", "left_unc", [0,  170,127], [0,0,-1], [1,0,0], stat, type, threshold);
runme(path_to_stats,path_to_edges,"left_cst", "left_cst", [0,  0,  255]  , [1,0,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges,"left_ilf", "left_ilf", [0,  255,0]  , [1,0,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges,"left_slf", "left_slf", [85, 170,127]  , [1,0,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges,"left_ifo", "left_ifo", [170,255,0]  , [1,0,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges, "right_unc", "right_unc", [0,  170,127], [0,0,-1], [-1,0,0], stat, type, threshold);
runme(path_to_stats,path_to_edges,"right_cst", "right_cst", [0,  0,  255]  , [-1,0,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges,"right_ilf", "right_ilf", [0,  255,0]  , [-1,0,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges,"right_slf", "right_slf", [85, 170,127]  , [-1,0,0], [0,0,1], stat, type, threshold);
runme(path_to_stats,path_to_edges,"right_ifo", "right_ifo", [170,255,0]  , [-1,0,0], [0,0,1], stat, type,  threshold);