#include <iostream>
#include <list>

#include <TROOT.h>
#include <TTree.h>
#include <TH1.h>
#include <TH1D.h>
#include <TProfile.h>
#include <TCanvas.h>

std::list<TH1*> object_list;

void register_object(TH1 *obj, char *name);
void unregister_objects();
void draw_beam();

double inner_wedge(double x, int side=1)
{
   x = (side==1)? -x : x;
   if (x > 0.25 && x < 0.813) {
      return 1.;
   }
   else if (x > 0.25 && x < 1.5) {
      return 0.1;
   }
   else if (x > 0.25 && x < 2.1) {
      return (side==1)? 0.40: 0.40;
   }
   else if (x > 0.25 && x < 2.5) {
      return (side==1)? 0.18 : 0.80;
   }
   else if (x > 0.25 && x < 3.4) {
      return (side==1)? 0.18 : 0.;
   }
   else if (x > 0.25 && x < 3.8) {
      return (side==1)? 0. : 0.;
   }
   return 0;
}

void unfold_inner(int side=1, double sigma=0.10)
{
   unregister_objects();

   // create a histogram representing the radial form of the inner wedge
   TH1D *form = new TH1D("innerform","inner form",512,-10,10);
   for (int i=1; i<=512; ++i) {
      double x = form->GetXaxis()->GetBinCenter(i);
      form->SetBinContent(i,inner_wedge(x,side));
   }
   register_object(form,"innerform");

   // take the FFT (real and imaginary parts) of the wedge form
   TH1 *formFFTreal = form->FFT(0,"RE R2C M");
   register_object(formFFTreal,"formFFTreal");
   TH1 *formFFTimag = form->FFT(0,"IM R2C M");
   register_object(formFFTimag,"formFFTimag");
   TH1 *formFFTmag = form->FFT(0,"MAG R2C M");
   register_object(formFFTmag,"formFFTmag");

   // take the FFT (real and imaginary parts) of the measured current profile
   TProfile *prof = new TProfile("prof","prof",512,-10,10,0,2);
   register_object(prof,"prof");
   TTree *run26_8e = (TTree*)gROOT->FindObject("run26_8e");
   if (side == 1) {
      run26_8e->Project("prof","ampcur[1]:xcm-14.3");
   }
   else {
      run26_8e->Project("prof","ampcur[2]:xcm-14.3");
   }
   TH1 *profFFTreal = prof->FFT(0,"RE R2C M");
   register_object(profFFTreal,"profFFTreal");
   TH1 *profFFTimag = prof->FFT(0,"IM R2C M");
   register_object(profFFTimag,"profFFTimag");

   // deconvolve the measured profile with the wedge form in Fourier space
   TH1 *workFFTr1 = prof->FFT(0,"RE R2C M");
   register_object(workFFTr1,"workFFTr1");
   workFFTr1->Multiply(formFFTreal);
   TH1 *workFFTi1 = prof->FFT(0,"IM R2C M");
   register_object(workFFTi1,"workFFTi1");
   workFFTi1->Multiply(formFFTimag);
   workFFTr1->Add(workFFTi1,1);
   workFFTr1->Divide(formFFTmag);
   workFFTr1->Divide(formFFTmag);
   TH1 *workFFTr2 = prof->FFT(0,"RE R2C M");
   register_object(workFFTr2,"workFFTr2");
   workFFTr2->Multiply(formFFTimag);
   TH1 *workFFTi2 = prof->FFT(0,"IM R2C M");
   register_object(workFFTi2,"workFFTi2");
   workFFTi2->Multiply(formFFTreal);
   workFFTi2->Add(workFFTr2,-1);
   workFFTi2->Divide(formFFTmag);
   workFFTi2->Divide(formFFTmag);

   // apply a frequency-dependent filter to the deconvolved shape
   double n0 = 20/(6.28*sigma);
   for (int n=0; n<512; ++n) {
      double n1 = (n < 265)? n : 512-n;
      double attenfactor = exp(-0.5*(n1*n1)/(n0*n0));
      if (n > 20 && 514-n > 20) {
         attenfactor = 0;
      }
      workFFTr1->SetBinContent(n+1,workFFTr1->GetBinContent(n+1)*attenfactor);
      workFFTi2->SetBinContent(n+1,workFFTi2->GetBinContent(n+1)*attenfactor);
   }

   // transform the deconvolved shape back to coordinate space
   TH1 *beamreal = workFFTr1->FFT(0,"RE R2C M");
   register_object(beamreal,"beamreal");
   TH1 *beamimag = workFFTi2->FFT(0,"IM R2C M");
   register_object(beamimag,"beamimag");
   beamreal->Add(beamimag);
   beamreal->Scale(1./512.);

   // save the deconvolved shape as the beam profile centered on zero
   TH1D *beam = new TH1D("beam","beam profile",512,-10,10);
   register_object(beam,"beam");
   for (int i=1; i<=512; ++i) {
      beam->SetBinContent(1+(i+255)%512,beamreal->GetBinContent(i));
   }

   // for reference, convolve the deconvolved shape with the wedge form
   // for comparison with the original measured profile
   TH1D *workFFTr3 = new TH1D(*(TH1D*)formFFTreal);
   register_object(workFFTr3,"workFFTr3");
   TH1D *workFFTi3 = new TH1D(*(TH1D*)formFFTimag);
   register_object(workFFTi3,"workFFTi3");
   workFFTr3->Multiply(workFFTr1);
   workFFTi3->Multiply(workFFTi2);
   workFFTr3->Add(workFFTi3,-1);
   TH1D *workFFTr4 = new TH1D(*(TH1D*)formFFTreal);
   register_object(workFFTr4,"workFFTr4");
   TH1D *workFFTi4 = new TH1D(*(TH1D*)formFFTimag);
   register_object(workFFTi4,"workFFTi4");
   workFFTr4->Multiply(workFFTi2);
   workFFTi4->Multiply(workFFTr1);
   workFFTi4->Add(workFFTr4,1);
   TH1 *recprofreal = workFFTr3->FFT(0,"RE R2C M");
   register_object(recprofreal,"recprofreal");
   TH1 *recprofimag = workFFTi4->FFT(0,"IM R2C M");
   register_object(recprofimag,"recprofimag");
   recprofreal->Add(recprofimag,1);
   recprofreal->Scale(1./512.);

   // copy the reconstructed profile into a histogram with correct x values
   TH1D *recprof = new TH1D("recprof","reconstructed profile",512,-10.,10.);
   register_object(recprof,"recprof");
   for (int i=1; i<=512; ++i) {
      recprof->SetBinContent(i,recprofreal->GetBinContent(i));
   }

   draw_beam();
}

void register_object(TH1 *obj, char *name)
{
   obj->SetName(name);
   object_list.push_back(obj);
}

void unregister_objects()
{
   std::list<TH1*>::iterator iter;
   for (iter = object_list.begin();
        iter != object_list.end();
        ++iter)
   {
      delete *iter;
   }
   object_list.clear();
}

void draw_beam()
{
   TCanvas *c1 = (TCanvas*)gROOT->FindObject("c1");
   if (c1 == 0) {
      c1 = new TCanvas("c1","c1",5,5,560,560);
   }
   TH1 *beam = (TH1*)gROOT->FindObject("beam");
   if (beam) {
      beam->Draw();
   }
}