#include <iostream>
#include <iomanip>
#include <cmath>

#include <TROOT.h>

#include "ICOMETLog.hxx"
#include "IReconState.hxx"

ClassImp(COMET::IReconState);

COMET::IReconState::IReconState() { }

COMET::IReconState::IReconState(const IReconState& state) 
    : TObject(state), fValues(state.fValues), fFieldNames(state.fFieldNames) { }

COMET::IReconState::~IReconState() { }

std::string COMET::IReconState::GetStateFields(void) const {
    // Construct a type name out of the field names.  This in turn is used by
    // the ICorrValues class to construct a type hash which is used to make
    // sure that operations are done on compatible ICorrValues objects.
    std::string typeName;
    for (std::vector<std::string>::const_iterator n = fFieldNames.begin();
         n != fFieldNames.end();
         ++n) {
        typeName += *n;
        typeName += " ";
    };
    return typeName;
}

// Build the internal state vector.
void COMET::IReconState::Init() {
    fValues.ResizeTo(fFieldNames.size());
    fValues.SetType(GetStateFields().c_str());
}

int COMET::IReconState::GetDimensions() const {
    return fValues.GetDimensions();
}

double COMET::IReconState::GetValue(int i) const {
    return fValues.GetValue(i);
}

void COMET::IReconState::SetValue(int i, double val) {
    return fValues.SetValue(i, val);
}

double COMET::IReconState::GetCovarianceValue(int i, int j) const {
    return fValues.GetCovarianceValue(i,j);
}

void COMET::IReconState::SetCovarianceValue(int i, int j, double val) {
    fValues.SetCovarianceValue(i,j,val);
}

void COMET::IReconState::SetFree(int i) {
    fValues.SetFree(i);
}

bool COMET::IReconState::IsFree(int i) const {
    return fValues.IsFree(i);
}

bool COMET::IReconState::IsFree(double v) const {
    return fValues.IsFree(v);
}

void COMET::IReconState::SetFixed(int i) {
    fValues.SetFixed(i);
}

bool COMET::IReconState::IsFixed(int i) const {
    return fValues.IsFixed(i);
}

bool COMET::IReconState::IsFixed(double v) const {
    return fValues.IsFixed(v);
}

void COMET::IReconState::Validate() {
    fValues.Validate(true);
}

COMET::ICorrValues COMET::IReconState::ProjectState(const COMET::IHandle<COMET::IReconState>& state) {
    return state->fValues;
}

/// Print the object information.
void COMET::IReconState::ls(Option_t*) const {
    TROOT::IndentLevel();
    std::cout << ClassName() << "(" << this << ")::" << std::endl;
    TROOT::IncreaseDirLevel();
    TROOT::IndentLevel();
    std::cout << GetStateFields() << std::endl;
    TROOT::IncreaseDirLevel();
    std::ios::fmtflags save = std::cout.flags();
    for (int i = 0; i<GetDimensions(); ++i) {
        if (IsFree(i)) continue;
        TROOT::IndentLevel();
        std::cout << "  " << std::setw(6) << fFieldNames[i];
        std::cout << ":: " 
                  << std::setw(9) << std::setprecision(3) 
                  << GetValue(i);
        if (IsFixed(i)) std::cout << "    fixed";
        else {
            std::cout << " +- " 
                      << std::setw(5) << std::setprecision(3) 
                      << std::sqrt(GetCovarianceValue(i,i));
        }
        std::cout << std::endl;
    }
    std::cout.flags(save);
    TROOT::DecreaseDirLevel();
    TROOT::DecreaseDirLevel();
}

////////////////////////////////////////////////////////
COMET::IMReconState::~IMReconState() {}

COMET::IMReconState::IMReconState(): fThis(NULL) {}

COMET::IMReconState::IMReconState(IReconState* state): fThis(state) {}

COMET::IReconState& COMET::IMReconState::GetThis() const {
    if (!fThis) {
        COMETError("Base pointer (fThis) is NULL");
    }
    return *fThis;
}

///////////////////////////////////////////////////////
COMET::IMEDepositState::~IMEDepositState() {}

COMET::IMEDepositState::IMEDepositState() {
    fEDepositIndex = fLocalNames.size();
    fLocalNames.push_back("EDeposit");
}

double COMET::IMEDepositState::GetEDeposit() const {
    return GetThis().GetValue(fEDepositIndex);
}

void COMET::IMEDepositState::SetEDeposit(double enr) {
    GetThis().SetValue(fEDepositIndex,enr);
}

double COMET::IMEDepositState::GetEDepositVariance() const {
    return GetThis().GetCovarianceValue(fEDepositIndex,fEDepositIndex);
}

void COMET::IMEDepositState::SetEDepositVariance(double var) {
    GetThis().SetCovarianceValue(fEDepositIndex,fEDepositIndex,var);
}

COMET::ICorrValues COMET::IMEDepositState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMEDepositState::GetSize());
    values.SetType("EDeposit ");
    const IMEDepositState* eDepositState 
        = dynamic_cast<const IMEDepositState*>(GetPointer(proj));
    int base = 0;
    if (eDepositState) {
        const int offset = eDepositState->GetEDepositIndex();
        for (int i = 0; i < IMEDepositState::GetSize(); ++i) {
            values.SetValue(i+base,
                            eDepositState->
                            GetThis().GetValue(i+offset));
            for (int j = 0; j < IMEDepositState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    eDepositState->
                    GetThis().GetCovarianceValue(i+offset,
                                                         j+offset));
            }
        }
    }
    return values;
}

///////////////////////////////////////////////////////
COMET::IMPositionState::~IMPositionState() {}

COMET::IMPositionState::IMPositionState() {
    fPositionIndex = fLocalNames.size();
    fLocalNames.push_back("X");
    fLocalNames.push_back("Y");
    fLocalNames.push_back("Z");
    fLocalNames.push_back("T");
}

TLorentzVector COMET::IMPositionState::GetPosition() const {
    return TLorentzVector(GetThis().GetValue(fPositionIndex),
                          GetThis().GetValue(fPositionIndex+1),
                          GetThis().GetValue(fPositionIndex+2),
                          GetThis().GetValue(fPositionIndex+3));
}

void COMET::IMPositionState::SetPosition(double x, double y,
                                      double z, double t) {
    GetThis().SetValue(fPositionIndex,x);
    GetThis().SetValue(fPositionIndex+1,y);
    GetThis().SetValue(fPositionIndex+2,z);
    GetThis().SetValue(fPositionIndex+3,t);
}

void COMET::IMPositionState::SetPosition(const TLorentzVector& pos) {
    SetPosition(pos.X(), pos.Y(), pos.Z(), pos.T());
}

TLorentzVector COMET::IMPositionState::GetPositionVariance() const {
    return TLorentzVector(
        GetThis().GetCovarianceValue(fPositionIndex,
                                             fPositionIndex),
        GetThis().GetCovarianceValue(fPositionIndex+1,
                                             fPositionIndex+1),
        GetThis().GetCovarianceValue(fPositionIndex+2,
                                             fPositionIndex+2),
        GetThis().GetCovarianceValue(fPositionIndex+3,
                                             fPositionIndex+3));
}

void COMET::IMPositionState::SetPositionVariance(double x, double y, 
                                              double z, double t) {
    GetThis().SetCovarianceValue(fPositionIndex,
                               fPositionIndex,x);
    GetThis().SetCovarianceValue(fPositionIndex+1,
                               fPositionIndex+1,y);
    GetThis().SetCovarianceValue(fPositionIndex+2,
                               fPositionIndex+2,z);
    GetThis().SetCovarianceValue(fPositionIndex+3,
                               fPositionIndex+3,t);
}

void COMET::IMPositionState::SetPositionCovariance(int i, int j, double v) {
    GetThis().SetCovarianceValue(fPositionIndex+i,fPositionIndex+j,v);
}

double COMET::IMPositionState::GetPositionCovariance(int i, int j) {
    return GetThis().GetCovarianceValue(fPositionIndex+i,fPositionIndex+j);
}

COMET::ICorrValues
COMET::IMPositionState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMPositionState::GetSize());
    values.SetType("X Y Z T ");
    const IMPositionState* posState 
        = dynamic_cast<const IMPositionState*>(GetPointer(proj));
    int base = 0;
    if (posState) {
        const int offset = posState->GetPositionIndex();
        for (int i = 0; i < IMPositionState::GetSize(); ++i) {
            values.SetValue(i+base, 
                            posState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMPositionState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    posState->GetThis().GetCovarianceValue(i+offset,
                                                                   j+offset));
            }
        }
    }
    return values;
}

///////////////////////////////////////////////////////
COMET::IMDirectionState::~IMDirectionState() {}

COMET::IMDirectionState::IMDirectionState() {
    fDirectionIndex = fLocalNames.size();
    fLocalNames.push_back("DX");
    fLocalNames.push_back("DY");
    fLocalNames.push_back("DZ");
}

TVector3 COMET::IMDirectionState::GetDirection() const {
    return TVector3(GetThis().GetValue(fDirectionIndex),
                    GetThis().GetValue(fDirectionIndex+1),
                    GetThis().GetValue(fDirectionIndex+2));
}

void COMET::IMDirectionState::SetDirection(double x, double y,
                                      double z) {
    GetThis().SetValue(fDirectionIndex,x);
    GetThis().SetValue(fDirectionIndex+1,y);
    GetThis().SetValue(fDirectionIndex+2,z);
}

void COMET::IMDirectionState::SetDirection(const TVector3& dir) {
    SetDirection(dir.X(), dir.Y(), dir.Z());
}

TVector3 COMET::IMDirectionState::GetDirectionVariance() const {
    return TVector3(GetThis().GetCovarianceValue(fDirectionIndex,
                                                    fDirectionIndex),
                    GetThis().GetCovarianceValue(fDirectionIndex+1,
                                                    fDirectionIndex+1),
                    GetThis().GetCovarianceValue(fDirectionIndex+2,
                                                    fDirectionIndex+2));
}

void COMET::IMDirectionState::SetDirectionVariance(double x, double y, 
                                                double z) {
    GetThis().SetCovarianceValue(fDirectionIndex,
                                    fDirectionIndex,x);
    GetThis().SetCovarianceValue(fDirectionIndex+1,
                                    fDirectionIndex+1,y);
    GetThis().SetCovarianceValue(fDirectionIndex+2,
                                    fDirectionIndex+2,z);
}

void COMET::IMDirectionState::SetDirectionCovariance(int i, int j, double v) {
    GetThis().SetCovarianceValue(fDirectionIndex+i,fDirectionIndex+j,v);
}

double COMET::IMDirectionState::GetDirectionCovariance(int i, int j) {
    return GetThis().GetCovarianceValue(fDirectionIndex+i,fDirectionIndex+j);
}

COMET::ICorrValues 
COMET::IMDirectionState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMDirectionState::GetSize());
    values.SetType("DX DY DZ ");
    const IMDirectionState* dirState 
        = dynamic_cast<const IMDirectionState*>(GetPointer(proj));
    int base = 0;
    if (dirState) {
        const int offset = dirState->GetDirectionIndex();
        for (int i = 0; i < IMDirectionState::GetSize(); ++i) {
            values.SetValue(i+base, 
                            dirState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMDirectionState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    dirState->GetThis().GetCovarianceValue(i+offset,
                                                                   j+offset));
            }
        }
    }
    return values;
}

///////////////////////////////////////////////////////
COMET::IMMomentumState::~IMMomentumState() {}

COMET::IMMomentumState::IMMomentumState() {
    fMomentumIndex = fLocalNames.size();
    fLocalNames.push_back("Momentum");
}

double COMET::IMMomentumState::GetMomentum() const {
    return GetThis().GetValue(fMomentumIndex);
}

void COMET::IMMomentumState::SetMomentum(double enr) {
    GetThis().SetValue(fMomentumIndex,enr);
}

double COMET::IMMomentumState::GetMomentumVariance() const {
    return GetThis().GetCovarianceValue(fMomentumIndex,fMomentumIndex);
}

void COMET::IMMomentumState::SetMomentumVariance(double var) {
    GetThis().SetCovarianceValue(fMomentumIndex,fMomentumIndex,var);
}

COMET::ICorrValues COMET::IMMomentumState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMMomentumState::GetSize());
    values.SetType("Momentum ");
    const IMMomentumState* momState 
        = dynamic_cast<const IMMomentumState*>(GetPointer(proj));
    int base = 0;
    if (momState) {
        const int offset = momState->GetMomentumIndex();
        for (int i = 0; i < IMMomentumState::GetSize(); ++i) {
            values.SetValue(i+base, 
                            momState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMMomentumState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    momState->GetThis().GetCovarianceValue(i+offset,
                                                                   j+offset));
            }
        }
    }
    return values;
}

///////////////////////////////////////////////////////
COMET::IMChargeState::~IMChargeState() {}

COMET::IMChargeState::IMChargeState() {
    fChargeIndex = fLocalNames.size();
    fLocalNames.push_back("Charge");
}

double COMET::IMChargeState::GetCharge() const {
    return GetThis().GetValue(fChargeIndex);
}

void COMET::IMChargeState::SetCharge(double enr) {
    GetThis().SetValue(fChargeIndex,enr);
}

double COMET::IMChargeState::GetChargeVariance() const {
    return GetThis().GetCovarianceValue(fChargeIndex,fChargeIndex);
}

void COMET::IMChargeState::SetChargeVariance(double var) {
    GetThis().SetCovarianceValue(fChargeIndex,fChargeIndex,var);
}

COMET::ICorrValues COMET::IMChargeState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMChargeState::GetSize());
    values.SetType("Charge ");
    const IMChargeState* chgState 
        = dynamic_cast<const IMChargeState*>(GetPointer(proj));
    int base = 0;
    if (chgState) {
        const int offset = chgState->GetChargeIndex();
        for (int i = 0; i < IMChargeState::GetSize(); ++i) {
            values.SetValue(i+base,
                            chgState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMChargeState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    chgState->GetThis().GetCovarianceValue(i+offset,
                                                                   j+offset));
            }
        }
    }
    return values;
}

///////////////////////////////////////////////////////
COMET::IMCurvatureState::~IMCurvatureState() {}

COMET::IMCurvatureState::IMCurvatureState() {
    fCurvatureIndex = fLocalNames.size();
    fLocalNames.push_back("Curvature");
}

double COMET::IMCurvatureState::GetCurvature() const {
    return GetThis().GetValue(fCurvatureIndex);
}

void COMET::IMCurvatureState::SetCurvature(double enr) {
    GetThis().SetValue(fCurvatureIndex,enr);
}

double COMET::IMCurvatureState::GetCurvatureVariance() const {
    return GetThis().GetCovarianceValue(fCurvatureIndex,
                                                fCurvatureIndex);
}

void COMET::IMCurvatureState::SetCurvatureVariance(double var) {
    GetThis().SetCovarianceValue(fCurvatureIndex,fCurvatureIndex,var);
}

COMET::ICorrValues COMET::IMCurvatureState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMCurvatureState::GetSize());
    values.SetType("Curvature ");
    const IMCurvatureState* curvState 
        = dynamic_cast<const IMCurvatureState*>(GetPointer(proj));
    int base = 0;
    if (curvState) {
        const int offset = curvState->GetCurvatureIndex();
        for (int i = 0; i < IMCurvatureState::GetSize(); ++i) {
            values.SetValue(i+base,
                            curvState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMCurvatureState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    curvState->GetThis().GetCovarianceValue(i+offset,
                                                                    j+offset));
            }
        }
    }
    return values;
}

///////////////////////////////////////////////////////
COMET::IMConeState::~IMConeState() {}

COMET::IMConeState::IMConeState() {
    fConeIndex = fLocalNames.size();
    fLocalNames.push_back("C1");
    fLocalNames.push_back("C2");
}

TVector3 COMET::IMConeState::GetCone() const {
    return TVector3(GetThis().GetValue(fConeIndex),
                    GetThis().GetValue(fConeIndex+1),
                    0);
}

void COMET::IMConeState::SetCone(double x, double y) {
    GetThis().SetValue(fConeIndex,x);
    GetThis().SetValue(fConeIndex+1,y);
}

TVector3 COMET::IMConeState::GetConeVariance() const {
    return TVector3(GetThis().GetCovarianceValue(fConeIndex,
                                                    fConeIndex),
                    GetThis().GetCovarianceValue(fConeIndex+1,
                                                    fConeIndex+1),
                    0);
}

void COMET::IMConeState::SetConeVariance(double x, double y) {
    GetThis().SetCovarianceValue(fConeIndex,
                                         fConeIndex,x);
    GetThis().SetCovarianceValue(fConeIndex+1,
                                         fConeIndex+1,y);
}

COMET::ICorrValues COMET::IMConeState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMConeState::GetSize());
    values.SetType("C1 C2 ");
    const IMConeState* coneState 
        = dynamic_cast<const IMConeState*>(GetPointer(proj));
    int base = 0;
    if (coneState) {
        const int offset = coneState->GetConeIndex();
        for (int i = 0; i < IMConeState::GetSize(); ++i) {
            values.SetValue(i+base,
                            coneState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMConeState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    coneState->GetThis().GetCovarianceValue(i+offset,
                                                                    j+offset));
            }
        }
    }
    return values;
}

///////////////////////////////////////////////////////
COMET::IMWidthState::~IMWidthState() {}

COMET::IMWidthState::IMWidthState() {
    fWidthIndex = fLocalNames.size();
    fLocalNames.push_back("W1");
    fLocalNames.push_back("W2");
}


TVector3 COMET::IMWidthState::GetWidth() const {
    return TVector3(GetThis().GetValue(fWidthIndex),
                    GetThis().GetValue(fWidthIndex+1),
                    0);
}

void COMET::IMWidthState::SetWidth(double x, double y) {
    GetThis().SetValue(fWidthIndex,x);
    GetThis().SetValue(fWidthIndex+1,y);
}

TVector3 COMET::IMWidthState::GetWidthVariance() const {
    return TVector3(GetThis().GetCovarianceValue(fWidthIndex,
                                                    fWidthIndex),
                    GetThis().GetCovarianceValue(fWidthIndex+1,
                                                    fWidthIndex+1),
                    0);
}

void COMET::IMWidthState::SetWidthVariance(double x, double y) {
    GetThis().SetCovarianceValue(fWidthIndex,
                                         fWidthIndex,x);
    GetThis().SetCovarianceValue(fWidthIndex+1,
                                         fWidthIndex+1,y);
}

COMET::ICorrValues COMET::IMWidthState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMWidthState::GetSize());
    values.SetType("W1 W2 ");
    const IMWidthState* wdthState 
        = dynamic_cast<const IMWidthState*>(GetPointer(proj));
    int base = 0;
    if (wdthState) {
        const int offset = wdthState->GetWidthIndex();
        for (int i = 0; i < IMWidthState::GetSize(); ++i) {
            values.SetValue(i+base,
                            wdthState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMWidthState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    wdthState->GetThis().GetCovarianceValue(i+offset,
                                                                    j+offset));
            }
        }
    }
    return values;
}

///////////////////////////////////////////////////////
COMET::IMPositionDirectionState::~IMPositionDirectionState() {}

COMET::ICorrValues COMET::IMPositionDirectionState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMPositionDirectionState::GetSize());
    values.SetType("X Y Z T DX DY DZ ");
    const IMPositionState* posState 
        = dynamic_cast<const IMPositionState*>(GetPointer(proj));
    int base = 0;
    if (posState) {
        int offset = posState->GetPositionIndex();
        for (int i = 0; i < IMPositionState::GetSize(); ++i) {
            values.SetValue(i+base, 
                            posState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMPositionState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    posState->GetThis().GetCovarianceValue(i+offset,
                                                                   j+offset));
            }
        }
    }
    base += IMDirectionState::GetSize();
    const IMDirectionState* dirState 
        = dynamic_cast<const IMDirectionState*>(GetPointer(proj));
    if (dirState) {
        int offset = dirState->GetDirectionIndex();
        for (int i = 0; i < IMDirectionState::GetSize(); ++i) {
            values.SetValue(i+base, 
                            dirState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMDirectionState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    dirState->GetThis().GetCovarianceValue(i+offset,
                                                                   j+offset));
            }
        }
    }
    return values;
}

///////////////////////////////////////////////////////
COMET::IMPosDirCurvState::~IMPosDirCurvState() {}

COMET::ICorrValues COMET::IMPosDirCurvState::ProjectState(const COMET::IHandle<COMET::IReconState>& proj) {
    ICorrValues values(IMPosDirCurvState::GetSize());
    values.SetType("X Y Z T DX DY DZ Curvature ");
    const IMPositionState* posState 
        = dynamic_cast<const IMPositionState*>(GetPointer(proj));
    int base = 0;
    if (posState) {
        int offset = posState->GetPositionIndex();
        for (int i = 0; i < IMPositionState::GetSize(); ++i) {
            values.SetValue(i+base, 
                            posState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMPositionState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    posState->GetThis().GetCovarianceValue(i+offset,
                                                                   j+offset));
            }
        }
    }
    base += IMDirectionState::GetSize();
    const IMDirectionState* dirState 
        = dynamic_cast<const IMDirectionState*>(GetPointer(proj));
    if (dirState) {
        int offset = dirState->GetDirectionIndex();
        for (int i = 0; i < IMDirectionState::GetSize(); ++i) {
            values.SetValue(i+base, 
                            dirState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMDirectionState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    dirState->GetThis().GetCovarianceValue(i+offset,
                                                                   j+offset));
            }
        }
    }
    base += IMDirectionState::GetSize();
    const IMCurvatureState* curvState 
        = dynamic_cast<const IMCurvatureState*>(GetPointer(proj));
    if (curvState) {
        int offset = curvState->GetCurvatureIndex();
        for (int i = 0; i < IMCurvatureState::GetSize(); ++i) {
            values.SetValue(i+base, 
                            curvState->GetThis().GetValue(i+offset));
            for (int j = 0; j < IMCurvatureState::GetSize(); ++j) {
                values.SetCovarianceValue(
                    i+base, j+base,
                    curvState->GetThis().GetCovarianceValue(i+offset,
                                                                    j+offset));
            }
        }
    }
    return values;
}