/////////////////////////////////////////////////////////////////////////////////
//
//  Copyright (C) 2017     Statoil ASA
// 
//  ResInsight is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
// 
//  ResInsight is distributed in the hope that it will be useful, but WITHOUT ANY
//  WARRANTY; without even the implied warranty of MERCHANTABILITY or
//  FITNESS FOR A PARTICULAR PURPOSE.
// 
//  See the GNU General Public License at <http://www.gnu.org/licenses/gpl.html> 
//  for more details.
//
/////////////////////////////////////////////////////////////////////////////////

#include "RimWellAllocationPlot.h"

#include "RiaApplication.h"

#include "RigEclipseCaseData.h"
#include "RigSimulationWellCenterLineCalculator.h"
#include "RigSimulationWellCoordsAndMD.h"
#include "RigSingleWellResultsData.h"

#include "RimEclipseCase.h"
#include "RimEclipseCellColors.h"
#include "RimEclipseResultCase.h"
#include "RimEclipseView.h"
#include "RimEclipseWell.h"
#include "RimEclipseWellCollection.h"
#include "RimFlowDiagSolution.h"
#include "RimTotalWellAllocationPlot.h"
#include "RimWellFlowRateCurve.h"
#include "RimWellLogPlot.h"
#include "RimWellLogTrack.h"

#include "RiuMainPlotWindow.h"
#include "RiuWellAllocationPlot.h"

CAF_PDM_SOURCE_INIT(RimWellAllocationPlot, "WellAllocationPlot");


//==================================================================================================
/// 
/// 
//==================================================================================================

class RigAccWellFlowCalculator
{

public:
    RigAccWellFlowCalculator(const std::vector< std::vector <cvf::Vec3d> >& pipeBranchesCLCoords,
                             const std::vector< std::vector <RigWellResultPoint> >& pipeBranchesCellIds):
        m_pipeBranchesCLCoords(pipeBranchesCLCoords), m_pipeBranchesCellIds(pipeBranchesCellIds)
    {
        m_accConnectionFlowPrBranch.resize(m_pipeBranchesCellIds.size());

        calculateAccumulatedFlowPrConnection(0,1);
    }

    // Returned pair is connection number from top of well with accumulated flow

    const std::vector<std::pair <size_t, double> >& accumulatedFlowPrConnection(size_t branchIdx)
    {
        return m_accConnectionFlowPrBranch[branchIdx];
    }

private:



    const std::vector<std::pair <size_t, double> >& calculateAccumulatedFlowPrConnection(size_t branchIdx, size_t startConnectionNumberFromTop)
    {
        std::vector<std::pair <size_t, double> >& accConnectionFlow = m_accConnectionFlowPrBranch[branchIdx];

        const std::vector<RigWellResultPoint>& branchCells =  m_pipeBranchesCellIds[branchIdx];
        int clSegIdx = static_cast<int>( branchCells.size()) - 1;
        double accFlow = 0.0;

        size_t prevGridIdx = -1;
        size_t prevGridCellIdx = -1;

        std::vector<size_t> resPointToConnectionIndexFromBottom;
        resPointToConnectionIndexFromBottom.resize(branchCells.size(), -1);

        size_t connIdxFromBottom = 0;

        while (clSegIdx >= 0)
        {
            resPointToConnectionIndexFromBottom[clSegIdx] = connIdxFromBottom;

            if ( branchCells[clSegIdx].m_gridIndex != prevGridIdx
                && branchCells[clSegIdx].m_gridCellIndex != prevGridIdx )
            {
                ++connIdxFromBottom;
            }

            prevGridIdx = branchCells[clSegIdx].m_gridIndex ;
            prevGridCellIdx = branchCells[clSegIdx].m_gridCellIndex;

            --clSegIdx;
        }

        size_t prevConnIndx = -1;
        clSegIdx = static_cast<int>( branchCells.size()) - 1;

        while (clSegIdx >= 0)
        {
            // Skip point if referring to the same cell as in the previous point did
            {
                if (resPointToConnectionIndexFromBottom[clSegIdx] == prevConnIndx)
                {
                    --clSegIdx;
                    continue;
                }

                prevConnIndx = resPointToConnectionIndexFromBottom[clSegIdx];
            }

            size_t connNumFromTop = connecionIdxFromTop(resPointToConnectionIndexFromBottom, clSegIdx) + startConnectionNumberFromTop;

            accFlow += branchCells[clSegIdx].m_flowRate;
            std::vector<size_t> downstreamBranches = findDownstreamBranchIdxs(branchCells[clSegIdx]);
            for (size_t dsBidx :  downstreamBranches )
            {
                if ( dsBidx != branchIdx &&  m_accConnectionFlowPrBranch[dsBidx].size() == 0) // Not this branch or already calculated
                {
                    const std::pair <size_t, double>& brancFlowPair = calculateAccumulatedFlowPrConnection(dsBidx, connNumFromTop).back();
                    accFlow += brancFlowPair.second;
                }
            }

            accConnectionFlow.push_back(std::make_pair(connNumFromTop, accFlow));

            --clSegIdx;
        }

        return m_accConnectionFlowPrBranch[branchIdx];
    }

    static size_t connecionIdxFromTop( std::vector<size_t> resPointToConnectionIndexFromBottom, size_t clSegIdx)
    {
        return resPointToConnectionIndexFromBottom.front() - resPointToConnectionIndexFromBottom[clSegIdx];
    }

    std::vector<size_t> findDownstreamBranchIdxs(const RigWellResultPoint& connectionPoint)
    {
        std::vector<size_t> downStreamBranchIdxs;

        for (size_t bIdx = 0; bIdx < m_pipeBranchesCellIds.size(); ++bIdx)
        {
            if ( m_pipeBranchesCellIds[bIdx][0].m_gridIndex == connectionPoint.m_gridIndex 
                && m_pipeBranchesCellIds[bIdx][0].m_gridCellIndex == connectionPoint.m_gridCellIndex)
            {
                downStreamBranchIdxs.push_back(bIdx);        
            }
        }
        return downStreamBranchIdxs;
    }


    const std::vector< std::vector <cvf::Vec3d> >& m_pipeBranchesCLCoords;
    const std::vector< std::vector <RigWellResultPoint> >& m_pipeBranchesCellIds;

    std::vector< std::vector<std::pair <size_t, double> > > m_accConnectionFlowPrBranch;

};


//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
RimWellAllocationPlot::RimWellAllocationPlot()
{
    CAF_PDM_InitObject("Well Allocation Plot", ":/newIcon16x16.png", "", "");

    CAF_PDM_InitField(&m_userName, "PlotDescription", QString("Flow Diagnostics Plot"), "Name", "", "", "");
    m_userName.uiCapability()->setUiReadOnly(true);

    CAF_PDM_InitField(&m_showPlotTitle, "ShowPlotTitle", true, "Show Plot Title", "", "", "");

    CAF_PDM_InitFieldNoDefault(&m_case, "CurveCase", "Case", "", "", "");
    m_case.uiCapability()->setUiTreeChildrenHidden(true);

    CAF_PDM_InitField(&m_timeStep, "PlotTimeStep", 0, "Time Step", "", "", "");
    CAF_PDM_InitField(&m_wellName, "WellName", QString("None"), "Well", "", "", "");
    CAF_PDM_InitFieldNoDefault(&m_flowDiagSolution, "FlowDiagSolution", "Flow Diag Solution", "", "", "");

    CAF_PDM_InitFieldNoDefault(&m_accumulatedWellFlowPlot, "AccumulatedWellFlowPlot", "Accumulated Well Flow", "", "", "");
    m_accumulatedWellFlowPlot.uiCapability()->setUiHidden(true);
    m_accumulatedWellFlowPlot = new RimWellLogPlot;

    CAF_PDM_InitFieldNoDefault(&m_totalWellAllocationPlot, "TotalWellFlowPlot", "Total Well Flow", "", "", "");
    m_totalWellAllocationPlot.uiCapability()->setUiHidden(true);

    m_totalWellAllocationPlot = new RimTotalWellAllocationPlot;

    this->setAsPlotMdiWindow();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
RimWellAllocationPlot::~RimWellAllocationPlot()
{
    removeMdiWindowFromMdiArea();
    
    delete m_accumulatedWellFlowPlot();
    delete m_totalWellAllocationPlot();

    deleteViewWidget();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
void RimWellAllocationPlot::setFromSimulationWell(RimEclipseWell* simWell)
{
    RimEclipseView* eclView;
    simWell->firstAncestorOrThisOfType(eclView);
    RimEclipseResultCase* eclCase;
    simWell->firstAncestorOrThisOfType(eclCase);

    m_case = eclCase;
    m_wellName = simWell->wellResults()->m_wellName;
    m_timeStep = eclView->currentTimeStep();

    m_flowDiagSolution = eclView->cellResult()->flowDiagSolution();

    updateFromWell();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
void RimWellAllocationPlot::deleteViewWidget()
{
    if (m_wellAllocationPlotWidget)
    {
        m_wellAllocationPlotWidget->deleteLater();
        m_wellAllocationPlotWidget = nullptr;
    }
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
void RimWellAllocationPlot::updateFromWell()
{
    if (!m_case) return;

    const RigSingleWellResultsData* wellResults = nullptr;
    wellResults = m_case->reservoirData()->findWellResult(m_wellName);

    if (!wellResults) return;

    
    const RigWellResultFrame* wellResultFrame = nullptr;
    std::vector< std::vector <cvf::Vec3d> > pipeBranchesCLCoords;
    std::vector< std::vector <RigWellResultPoint> > pipeBranchesCellIds;

    RigSimulationWellCenterLineCalculator::calculateWellPipeCenterlineFromWellFrame(m_case->reservoirData(),
                                                                                    wellResults,
                                                                                    m_timeStep,
                                                                                    true,
                                                                                    true,
                                                                                    pipeBranchesCLCoords,
                                                                                    pipeBranchesCellIds);

    accumulatedWellFlowPlot()->setDescription("Accumulated Well Flow (" + m_wellName + ")");

    RigAccWellFlowCalculator wfCalculator(pipeBranchesCLCoords, pipeBranchesCellIds);

    // Delete existing tracks
    {
        std::vector<RimWellLogTrack*> tracks;
        accumulatedWellFlowPlot()->descendantsIncludingThisOfType(tracks);

        for (RimWellLogTrack* t : tracks)
        {
            accumulatedWellFlowPlot()->removeTrack(t);
            delete t;
        }
    }

    CVF_ASSERT(accumulatedWellFlowPlot()->trackCount() == 0);

    size_t branchCount = pipeBranchesCLCoords.size();
    for (size_t brIdx = 0; brIdx < branchCount; ++brIdx)
    {
        RimWellLogTrack* plotTrack = new RimWellLogTrack();

        // TODO: Description is overwritten by RimWellLogPlot::updateTrackNames()
        // Add flag to control if this behavior
        plotTrack->setDescription(QString("Branch %1").arg(brIdx + 1));

        accumulatedWellFlowPlot()->addTrack(plotTrack);

        #if 0
        std::vector<cvf::Vec3d> curveCoords;
        std::vector<double> flowRate;

        {
            std::vector<cvf::Vec3d> branchCoords = pipeBranchesCLCoords[brIdx];
            std::vector<RigWellResultPoint> branchResultPoints = pipeBranchesCellIds[brIdx];

            RigWellResultPoint prevResultPoint;

            for (size_t i = 0; i < branchResultPoints.size(); i++)
            {
                const RigWellResultPoint& resultPoint = branchResultPoints[i];

                if (i > 0 && prevResultPoint.m_gridCellIndex != resultPoint.m_gridCellIndex)
                {
                    // Add an extra curve point when a cell transition is detected
                    curveCoords.push_back(branchCoords[i]);
                    flowRate.push_back(prevResultPoint.m_flowRate);
                }

                curveCoords.push_back(branchCoords[i]);
                flowRate.push_back(branchResultPoints[i].m_flowRate);

                prevResultPoint = resultPoint;
            }
        }

        RigSimulationWellCoordsAndMD helper(curveCoords);
        #endif

        const std::vector<std::pair <size_t, double> >& flowPrConnection =  wfCalculator.accumulatedFlowPrConnection(brIdx);

        std::vector<double> connNumbers;
        std::vector<double> accFlow;

        for (const std::pair<size_t, double> & flowPair: flowPrConnection)
        {
            connNumbers.push_back(flowPair.first);
            accFlow.push_back(flowPair.second);
        }

        RimWellFlowRateCurve* curve = new RimWellFlowRateCurve;
        curve->setFlowValues(connNumbers, accFlow);

        plotTrack->addCurve(curve);

        curve->loadDataAndUpdate();
    }

    setDescription("Well Allocation (" + m_wellName + ")");
    accumulatedWellFlowPlot()->updateConnectedEditors();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
QWidget* RimWellAllocationPlot::viewWidget()
{
    return m_wellAllocationPlotWidget;
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
void RimWellAllocationPlot::zoomAll()
{
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
RimWellLogPlot* RimWellAllocationPlot::accumulatedWellFlowPlot()
{
    return m_accumulatedWellFlowPlot();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
RimTotalWellAllocationPlot* RimWellAllocationPlot::totalWellFlowPlot()
{
    return m_totalWellAllocationPlot();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
QList<caf::PdmOptionItemInfo> RimWellAllocationPlot::calculateValueOptions(const caf::PdmFieldHandle* fieldNeedingOptions, bool * useOptionsOnly)
{
    QList<caf::PdmOptionItemInfo> options;

    if (fieldNeedingOptions == &m_wellName)
    {
        std::set<QString> sortedWellNames;
        if ( m_case && m_case->reservoirData() )
        {
            const cvf::Collection<RigSingleWellResultsData>& wellRes = m_case->reservoirData()->wellResults();

            for ( size_t wIdx = 0; wIdx < wellRes.size(); ++wIdx )
            {
                sortedWellNames.insert(wellRes[wIdx]->m_wellName);
            }
        }

        QIcon simWellIcon(":/Well.png");
        for ( const QString& wname: sortedWellNames )
        {
            options.push_back(caf::PdmOptionItemInfo(wname, wname, false, simWellIcon));
        }

        if (options.size() == 0)
        {
            options.push_front(caf::PdmOptionItemInfo("None", nullptr));
        }
    }
    else if (fieldNeedingOptions == &m_timeStep)
    {
        QStringList timeStepNames;

        if (m_case && m_case->reservoirData())
        {
            timeStepNames = m_case->timeStepStrings();
        }

        for (int i = 0; i < timeStepNames.size(); i++)
        {
            options.push_back(caf::PdmOptionItemInfo(timeStepNames[i], i));
        }
    }

    return options;
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
QString RimWellAllocationPlot::wellName() const
{
    return m_wellName();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
void RimWellAllocationPlot::removeFromMdiAreaAndDeleteViewWidget()
{
    removeMdiWindowFromMdiArea();
    deleteViewWidget();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
void RimWellAllocationPlot::fieldChangedByUi(const caf::PdmFieldHandle* changedField, const QVariant& oldValue, const QVariant& newValue)
{
    RimViewWindow::fieldChangedByUi(changedField, oldValue, newValue);

    if (changedField == &m_userName ||
        changedField == &m_showPlotTitle)
    {
        updateMdiWindowTitle();
    }
    else if (changedField == &m_wellName)
    {
        updateFromWell();
    }
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
QImage RimWellAllocationPlot::snapshotWindowContent()
{
    QImage image;

    // TODO

    return image;
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
void RimWellAllocationPlot::setDescription(const QString& description)
{
    m_userName = description;
    this->updateMdiWindowTitle();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
QString RimWellAllocationPlot::description() const
{
    return m_userName();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
void RimWellAllocationPlot::loadDataAndUpdate()
{
    updateMdiWindowVisibility();
    updateFromWell();
}

//--------------------------------------------------------------------------------------------------
/// 
//--------------------------------------------------------------------------------------------------
QWidget* RimWellAllocationPlot::createViewWidget(QWidget* mainWindowParent)
{
    m_wellAllocationPlotWidget = new RiuWellAllocationPlot(this, mainWindowParent);
    return m_wellAllocationPlotWidget;
}