# Imports x,y,z positions in a fits file as points in a seperate meshes with material settings controlled by image value
# Meshes each contain multiple data points of similar values
# Does everything from the fits file directly

import Blender as B

import Blender
from Blender import *
import math
from math import *
import numpy
from numpy import *
import pyfits
from pyfits import *
import bpy
from sys import stdout 


srange = 50.0							# For colour scaling, maximum pixel value above which everything is white
n = 1000.0								# No. meshes to create
S = 15.0									# For colour scaling, minimum pixel value below which nothing is displayed

deltaS = (srange-S)/n
Sf = S + deltaS

scn = Blender.Scene.GetCurrent()

me=B.Mesh.New('myMesh')				# Create initial mesh

ma = Blender.Material.New()		# Create initial material
ma.setMode('Halo')
a = float(S) / srange
ma.setAlpha(a)
ma.setAdd(a)
ma.rgbCol =1.0, 1.0, 1.0
ma.setHaloSize(1.0)

print 'Converting fits file to internal array...'
	
FitsFile = pyfits.open('Filename.fits')
image = FitsFile[0].data

z = image.shape[0]	# No. channels					
y = image.shape[1]	# No. x pixels
x = image.shape[2]	# No. y pixels

temp = array(transpose(nonzero(image)),float)	# Make an array of x,y,z of all non-zero values in the fits file
vals = array(image[nonzero(image)])					# Make a list of all those non-zero values

#print vals
l = len(vals)
print l, 'non-zero values found'

temp.resize(l,4)							# Enlarge the array to hold the pixel values

for i in range(0,l):
	temp[i,3] = vals[i]
		

print 'Sorting array by value'
temp = temp[temp[:,3].argsort(),:]	# Sorts the array by values in the 4th column. Got it from the net. Don't know how it works.

#print temp[0:10,:]

templ = temp.shape[0]
tempm = temp[templ-1,3]									# Maximum value since array has already been sorted

print 'Displaying data'

for i in range(0,templ):
	
	#print (float(i)/float(templ))*100.0			

	V = temp[i,3]											# x,y,z axes are more or less arbitrary, depends how you want Blender to display
	x = temp[i,2]											# stuff
	y = temp[i,1]
	z = temp[i,0]
		
	if float(V) >= S and float(V) < Sf:				# If S/N correct, add this point to the mesh		
	
		me.verts.extend(float(x),float(z),float(y))
											
	if float(V) >= Sf or float(V) == tempm:		# If S/N too high, set the materials of the mesh and create the object
		
		me.materials = [ma]
		for f in me.faces:
			f.mat = 0
			
		ob=scn.objects.new(me,'Gas')
		
		ma = Blender.Material.New()					# Setup a new material and mesh for the next set of data
		ma.setMode('Halo')
		
		if float(V) < srange :							# Checks for S/N cutoff
			a = float(V) / 50.0
		else:
			a = 1.0
			
		ma.setAlpha(a)
		ma.setAdd(a)
		ma.rgbCol =1.0, 1.0, 1.0
		ma.setHaloSize(1.0)
		
		me=B.Mesh.New('myMesh')
		
		S = S + deltaS
		Sf = Sf + deltaS
	
		if S >= srange :					# If above S/N cutoff, set the check to a huge value so only create 1 more object
												# and material			
			S = srange						
			Sf = tempm 		
			

B.Redraw()

print 'Done !'

FitsFile.close()