#!/usr/bin/env python

import itk
from itk import RTK as rtk
import os
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import scipy.optimize
import glob

# Script parameters
dir = "output/"
spacingVol = [ 0.5, 0.5, 0.5 ]
sizeOutputVol = [ 512, 1, 512 ]
order = 6;
refMu = 0.02;
originVol = [ -(sizeOutputVol[0]/2+0.5)*spacingVol[0], 0., -(sizeOutputVol[2]/2+0.5)*spacingVol[2]]

# List of filenames
fileNames = list()
for file in os.listdir(dir):
    if file.startswith("attenuation") and file.endswith(".mha"):
        fileNames.append(dir + file)

# Read in full geometry
geometryReader = rtk.ThreeDCircularProjectionGeometryXMLFileReader.New()
geometryReader.SetFilename ( "output/geometry.xml" )
geometryReader.GenerateOutputInformation()
geometry = geometryReader.GetGeometry();

ImageType = itk.Image[itk.F,3]
projReader = rtk.ProjectionsReader[ImageType].New()
projReader.SetFileNames(fileNames)

constantImageSource = rtk.ConstantImageSource[ImageType].New()
constantImageSource.SetOrigin( originVol )
constantImageSource.SetSpacing( spacingVol )
constantImageSource.SetSize( sizeOutputVol )
constantImageSource.SetConstant(0.)

# Create template
drawCyl = rtk.DrawEllipsoidImageFilter[ImageType, ImageType].New()
drawCyl.SetInput(constantImageSource.GetOutput())
drawCyl.SetDensity(refMu)
drawCyl.SetAngle(0)
drawCyl.SetAxis([100, 0, 100])
drawCyl.SetCenter([0, 0, 0])
drawCyl.Update()
template = drawCyl.GetOutput()
itk.imwrite(template, 'template.mha')
template = itk.GetArrayFromImage(template).flatten()

# Create weights (where the template should match)
drawCyl.SetDensity(1)
drawCyl.SetAngle(0)
drawCyl.SetAxis([125, 0, 125])
drawCyl.Update()
weights = drawCyl.GetOutput()
weights.DisconnectPipeline()

drawCyl.SetInput(weights)
drawCyl.SetDensity(-1)
drawCyl.SetAngle(0)
drawCyl.SetAxis([102, 0, 102])
drawCyl.Update()
weights = drawCyl.GetOutput()
weights.DisconnectPipeline()

drawCyl.SetInput(weights)
drawCyl.SetDensity(1)
drawCyl.SetAngle(0)
drawCyl.SetAxis([98, 0, 98])
weights = drawCyl.GetOutput()

itk.imwrite(weights, 'weights.mha')
weights = itk.GetArrayFromImage(weights).flatten()

# Create reconstructed images
waterPre = rtk.WaterPrecorrectionImageFilter[ImageType,ImageType].New()
wpcoeffs = np.zeros((order+1))

fdkFilter = rtk.FDKConeBeamReconstructionFilter[ImageType].New()
fdkFilter.SetGeometry( geometry )
fdkFilter.SetInput(0,constantImageSource.GetOutput())
fdkFilter.SetInput(1,waterPre.GetOutput())

fdks = [None] * (order+1)
for o in range(0,order+1):
    wpcoeffs[o-1] = 0.
    wpcoeffs[o] = 1.
    waterPre.SetCoefficients(wpcoeffs)
    waterPre.SetInput(projReader.GetOutput())
    fdkFilter.Update()
    fdk = fdkFilter.GetOutput()
    itk.imwrite(fdk, 'fdk%d.mha' % o )
    fdks[o] = itk.GetArrayFromImage(fdk).flatten()

# Create and solve the linear system of equation B.c= a to find the coeffs c 
a = np.zeros((order+1))
B = np.zeros((order+1,order+1))
for i in range(0,order+1):
    a[i] = np.sum(weights * fdks[i] * template)    
    for j in np.arange(i,order+1):
        B[i,j] = np.sum(weights * fdks[i] * fdks[j])
        B[j,i] = B[i,j]    
c = np.linalg.solve(B,a)

waterPre.SetCoefficients(c)
fdkFilter.Update()
itk.imwrite(fdkFilter.GetOutput(), 'fdk.mha' )

fdk = itk.imread('fdk.mha')
fdk1 = itk.imread('fdk1.mha')

plt.plot(itk.GetArrayFromImage(fdk)[:,0,256], label='Corrected')
plt.plot(itk.GetArrayFromImage(fdk1)[:,0,256], label='Uncorrected')
plt.legend()
plt.xlabel('Pixel number')
plt.ylabel('Attenuation')
plt.xlim(0,512)
plt.savefig('profile.png') 

