#include <TObject.h>
#include <TCollection.h> //for TIter
#include "IFieldPlotter.hxx"
#include "IFieldManager.hxx"
#include <ICOMETLog.hxx>
#include <IOADatabase.hxx>
#include <IGeomIdManager.hxx>
#include <sstream>
#include <cmath>

#include <TH2.h>
#include <TList.h>
#include <TRandom3.h>
#include <TGeoManager.h>
#include <TGeoMaterial.h>
#include <TGeoVolume.h>
#include <TGeoNode.h>

namespace {
    std::ostream& operator<<(std::ostream& stream,const TVector3& vect){
        stream<<"("<<vect.x()<<", "<<vect.y()<<", "<<vect.z() <<")";
        return stream;
    }
}

COMET::IFieldPlotter::IFieldPlotter():
    fUpVector(0,0,1),
    fNormal(0,1,0),
    fOrigin(0,0,0),
    fStepUp(1),
    fStepSide(1),
    fNSteps(500),
    fDirectory(TDirectory::CurrentDirectory()),
    fSampleGeometry(1e5)
    {
}

bool COMET::IFieldPlotter::Initialise(){
    COMETNamedLog("IFieldPlotter","Setting up with "
            "Normal: " <<fNormal<<
            " Up vector: "<< fUpVector
            );

    fDirOut=fNormal.Unit();
    fDirUp=fUpVector-fUpVector.Dot(fNormal)*fNormal;
    if(fDirUp.Mag()<1e-14) {
        COMETError("Up vector and plane normals should not be the same");
        return false;
    }
    fDirUp=fDirUp.Unit();
    fDirSide=fDirUp.Cross(fDirOut);

    COMETNamedLog("IFieldPlotter","Plane defined by: "
            "Normal: " <<fDirOut<<
            " Up vector: "<< fDirUp<<
            " Side vector:"<<fDirSide
            );
    return true;
}

bool COMET::IFieldPlotter::PlotPlane(TDirectory* output){
    TDirectory* oldwd=TDirectory::CurrentDirectory();
    if(output) output->cd();
    else if(fDirectory) fDirectory->cd();

    COMET::IFieldManager* theGlobalField = COMET::IFieldManager::GetObject();

    const double origin_side=fOrigin.Dot(fDirSide);
    const double origin_up=fOrigin.Dot(fDirUp);
    const double limit_side=fNSteps*fStepSide;
    const double limit_up=fNSteps*fStepUp;
    std::stringstream x_axis; x_axis<<"Distance along "<<fDirSide;
    std::stringstream y_axis; y_axis<<"Distance along "<<fDirUp;

    std::stringstream title;
    title<<"Magnitude of field through "<<fOrigin 
        <<";"<<x_axis.str()<<";"<<y_axis.str();
    TH2* hMagnitude=new TH2F("hMagnitude",title.str().c_str(),
            2*fNSteps+1 , -limit_side+origin_side , limit_side+origin_side ,
            2*fNSteps+1 , -limit_up+origin_up , limit_up+origin_up);

    title.str(""); 
    title<<"Component of field parallel to "<<fDirSide<<" through "<<fOrigin
        <<";"<<x_axis.str()<<";"<<y_axis.str();
    TH2* hSideways=new TH2F("hSideways",title.str().c_str(),
            2*fNSteps+1 , -limit_side+origin_side , limit_side+origin_side ,
            2*fNSteps+1 , -limit_up+origin_up , limit_up+origin_up);

    title.str(""); 
    title<<"Component of field parallel to "<<fDirUp<<" through "<<fOrigin
        <<";"<<x_axis.str()<<";"<<y_axis.str();
    TH2* hUpwards=new TH2F("hUpwards",title.str().c_str(),
            2*fNSteps+1 , -limit_side+origin_side , limit_side+origin_side ,
            2*fNSteps+1 , -limit_up+origin_up , limit_up+origin_up);

    title.str(""); 
    title<<"Component of field parallel to "<<fDirOut<<" through "<<fOrigin
        <<";"<<x_axis.str()<<";"<<y_axis.str();
    TH2* hOutwards=new TH2F("hOutwards",title.str().c_str(),
            2*fNSteps+1 , -limit_side+origin_side , limit_side+origin_side ,
            2*fNSteps+1 , -limit_up+origin_up , limit_up+origin_up);

    TH2* hGeometry=NULL;
    if(fSampleGeometry!=0){
        if(gGeoManager){
            title.str("");
            title<<"Geometry density  through "<<fOrigin 
                <<";"<<x_axis.str()<<";"<<y_axis.str();
            hGeometry=new TH2F("hGeometry",title.str().c_str(),
                    2*fNSteps+1 , -limit_side+origin_side , limit_side+origin_side ,
                    2*fNSteps+1 , -limit_up+origin_up , limit_up+origin_up);
        }
    }

    double point[3];
    double field_vectors[6];
    for(int i_side = -fNSteps; i_side<=fNSteps; ++i_side){
        double dist_side=fStepSide*i_side;
        TVector3 side= fDirSide*dist_side;
        dist_side+=origin_side;
        for(int i_up = -fNSteps; i_up<=fNSteps; ++i_up){
            double dist_up=fStepUp*i_up;
            TVector3 position= fDirUp*dist_up;
            position+= side;
            position+= fOrigin;
            dist_up+=origin_up;

            point[0]=position.x();
            point[1]=position.y();
            point[2]=position.z();
            for(int j=0; j<6;++j) field_vectors[j]=0;

            theGlobalField->GetFieldValue(point,field_vectors);
            TVector3 field(field_vectors[0],field_vectors[1],field_vectors[2]);
            field*=1/unit::tesla;

            hMagnitude->Fill(dist_side,dist_up,field.Mag());
            hUpwards->Fill(dist_side,dist_up,field.Dot(fDirUp));
            hSideways->Fill(dist_side,dist_up,field.Dot(fDirSide));
            hOutwards->Fill(dist_side,dist_up,field.Dot(fDirOut));
            if(hGeometry){
                const double density=GetDensity(position);
                hGeometry->Fill(dist_side,dist_up,density);
            }
        }
    }

    oldwd->cd();
    return true;
}

bool COMET::IFieldPlotter::PlotPlanes(int n_planes,double plane_sep){
    double start_depth=(n_planes-1)*0.5*plane_sep;
    int start=(n_planes-1)*0.5;
    int stop=start + (n_planes+1)%2;
    TVector3 orig_origin=fOrigin;
    fOrigin-=start_depth*fDirOut;

    for(int i_plane=-start;i_plane<=stop; ++i_plane){
        // create a name for the plane
        std::string slice_name=Form("slice_%g_%g_%g",fOrigin.x(),fOrigin.y(),fOrigin.z());
        TDirectory* dir=fDirectory->mkdir(slice_name.c_str());

        COMETNamedLog("IFieldPlotter","Running on plane ("<<i_plane<<") centred at "<<fOrigin);

        PlotPlane(dir);
        fOrigin+=fDirOut*plane_sep;
    }
    fOrigin=orig_origin;
    return true;
}

void COMET::IFieldPlotter::SampleGeometryInPlane(TH2* hGeometry){
    // Check for a valid geometry
    if(! gGeoManager) {
        COMETError("No valid geometry at this point, so I cannot produce the geometry sampling plot.");
        if(hGeometry) delete hGeometry;
        return;
    }
    gGeoManager->PushPoint();

    // Find the maximum density for material used in the geometry
    double max_density=0;
    TIter nextMaterial(gGeoManager->GetListOfMaterials());
    TObject* object=0;
    while( (object = nextMaterial() ) ){
        const TGeoMaterial* material=(const TGeoMaterial*) object;
        if(material) {
            const double density=material->GetDensity();
            if(density>max_density) max_density=density;
        }
    }

    // Until we've sampled fSampleGeometry points, 
    const double range_side=fNSteps*fStepSide;
    const double range_up  =fNSteps*fStepUp;
    const double origin_side=fOrigin.Dot(fDirSide);
    const double origin_up=fOrigin.Dot(fDirUp);
    TRandom3 generator;
    int points_skipped=0;
    const int maxIteratrions=10*fSampleGeometry;
    for(int n_samples=0; n_samples<fSampleGeometry; /*increment n_samples in loop*/ ){
        if(n_samples+points_skipped > maxIteratrions){
            break;
        }

        // Choose a point
        const double side=generator.Uniform(-range_side,range_side);
        const double up=generator.Uniform(-range_up,range_up);
        const TVector3 point=fOrigin+side*fDirSide+up*fDirUp;
        const double density=GetDensity(point);

        if(density<0 || generator.Rndm()>density/max_density){
            ++points_skipped;
            continue;
        }

        hGeometry->Fill(side+origin_side,up+origin_up,1.);
        ++n_samples;
    }

    gGeoManager->PopPoint();
}

double COMET::IFieldPlotter::GetDensity(const TVector3& point){
        // Factors of 10 come from units in geometry being cm, not mm as used by
        // field
        const TGeoNode* node = COMET::IOADatabase::Get().GeomId().FindNode(point);
        if(!node){
            COMETDebug("No node, perhaps in the world volume?");
            return -1;
        }
        const TGeoVolume* volume =node->GetVolume();
        if(!volume){
            COMETError("Node without a volume !");
            return -1;
        }
        const TGeoMaterial* material =volume->GetMaterial();
        return material->GetDensity();
}