From ea793c15f3150344ba04ca6f696a0af013427982 Mon Sep 17 00:00:00 2001
From: Gaute Lindkvist <lindkvis@gmail.com>
Date: Thu, 5 Dec 2019 13:51:47 +0100
Subject: [PATCH] Fix ensemble selection if statistics curves are not shown

---
 .../Flow/RimWellPlotTools.cpp                 | 31 +++++----
 .../ProjectDataModel/Flow/RimWellRftPlot.cpp  | 65 ++++++++++++-------
 2 files changed, 55 insertions(+), 41 deletions(-)

diff --git a/ApplicationCode/ProjectDataModel/Flow/RimWellPlotTools.cpp b/ApplicationCode/ProjectDataModel/Flow/RimWellPlotTools.cpp
index dfacdd0b4b..8abddde135 100644
--- a/ApplicationCode/ProjectDataModel/Flow/RimWellPlotTools.cpp
+++ b/ApplicationCode/ProjectDataModel/Flow/RimWellPlotTools.cpp
@@ -765,30 +765,29 @@ std::set<RiaRftPltCurveDefinition> RimWellPlotTools::curveDefsFromTimesteps(
                 }
             }
         }
-        else if ( addr.sourceType() == RifDataSourceForRftPlt::ENSEMBLE_RFT )
+        else if ( addr.ensemble() )
         {
-            if ( addr.ensemble() )
+            // Add individual summary curves
+            for ( RimSummaryCase* summaryCase : addr.ensemble()->allSummaryCases() )
             {
-                for ( RimSummaryCase* summaryCase : addr.ensemble()->allSummaryCases() )
+                if ( summaryCase && summaryCase->rftReader() )
                 {
-                    if ( summaryCase && summaryCase->rftReader() )
-                    {
-                        RifDataSourceForRftPlt summaryAddr( RifDataSourceForRftPlt::SUMMARY_RFT,
-                                                            summaryCase,
-                                                            addr.ensemble() );
+                    RifDataSourceForRftPlt summaryAddr( RifDataSourceForRftPlt::SUMMARY_RFT, summaryCase, addr.ensemble() );
 
-                        std::set<QDateTime> timeSteps = summaryCase->rftReader()->availableTimeSteps(
-                            wellPathNameOrSimWellName );
-                        for ( const QDateTime& time : timeSteps )
+                    std::set<QDateTime> timeSteps = summaryCase->rftReader()->availableTimeSteps(
+                        wellPathNameOrSimWellName );
+                    for ( const QDateTime& time : timeSteps )
+                    {
+                        if ( selectedTimeStepSet.count( time ) )
                         {
-                            if ( selectedTimeStepSet.count( time ) )
-                            {
-                                curveDefs.insert(
-                                    RiaRftPltCurveDefinition( summaryAddr, wellPathNameOrSimWellName, time ) );
-                            }
+                            curveDefs.insert( RiaRftPltCurveDefinition( summaryAddr, wellPathNameOrSimWellName, time ) );
                         }
                     }
                 }
+            }
+            // Add statistics curves
+            if ( addr.sourceType() == RifDataSourceForRftPlt::ENSEMBLE_RFT )
+            {
                 std::set<QDateTime> statTimeSteps = addr.ensemble()->rftTimeStepsForWell( wellPathNameOrSimWellName );
                 for ( const QDateTime& time : statTimeSteps )
                 {
diff --git a/ApplicationCode/ProjectDataModel/Flow/RimWellRftPlot.cpp b/ApplicationCode/ProjectDataModel/Flow/RimWellRftPlot.cpp
index 2411a47a6a..5ffd83ddfa 100644
--- a/ApplicationCode/ProjectDataModel/Flow/RimWellRftPlot.cpp
+++ b/ApplicationCode/ProjectDataModel/Flow/RimWellRftPlot.cpp
@@ -370,6 +370,11 @@ void RimWellRftPlot::updateEditorsFromCurves()
     {
         if ( curveDef.address().sourceType() == RifDataSourceForRftPlt::OBSERVED )
             selectedSources.insert( RifDataSourceForRftPlt( RifDataSourceForRftPlt::OBSERVED ) );
+        else if ( curveDef.address().sourceType() == RifDataSourceForRftPlt::SUMMARY_RFT )
+        {
+            selectedSources.insert(
+                RifDataSourceForRftPlt( RifDataSourceForRftPlt::ENSEMBLE_RFT, curveDef.address().ensemble() ) );
+        }
         else
             selectedSources.insert( curveDef.address() );
 
@@ -1236,6 +1241,33 @@ void RimWellRftPlot::defineCurveColorsAndSymbols( const std::set<RiaRftPltCurveD
         }
     }
 
+    auto ensembles = selectedEnsembles();
+    for ( RimWellRftEnsembleCurveSet* curveSet : m_ensembleCurveSets() )
+    {
+        CAF_ASSERT( curveSet );
+        auto ensemble_it = std::find_if( ensembles.begin(),
+                                         ensembles.end(),
+                                         [&curveSet]( const RimSummaryCaseCollection* ensemble ) {
+                                             return curveSet->ensemble() == ensemble;
+                                         } );
+        if ( ensemble_it != ensembles.end() )
+        {
+            curveSet->initializeLegend();
+
+            if ( curveSet->legendConfig()->showLegend() && curveSet->colorMode() == ColorMode::BY_ENSEMBLE_PARAM &&
+                 !curveSet->currentEnsembleParameter().isEmpty() )
+            {
+                if ( !m_ensembleLegendFrames[curveSet] )
+                {
+                    m_ensembleLegendFrames[curveSet] = new RiuCvfOverlayItemWidget( track->viewer(),
+                                                                                    track->viewer()->canvas() );
+                }
+                m_ensembleLegendFrames[curveSet]->updateFromOverlayItem( curveSet->legendConfig()->titledOverlayFrame() );
+                track->viewer()->addOverlayFrame( m_ensembleLegendFrames[curveSet] );
+            }
+        }
+    }
+
     std::vector<cvf::Color3f> colorTable;
     RiaColorTables::summaryCurveDefaultPaletteColors().color3fArray().toStdVector( &colorTable );
 
@@ -1252,34 +1284,17 @@ void RimWellRftPlot::defineCurveColorsAndSymbols( const std::set<RiaRftPltCurveD
         auto colorTableIndex  = m_dataSourceColors.size();
         auto symbolTableIndex = m_timeStepSymbols.size();
 
-        RifDataSourceForRftPlt address = curveDefToAdd.address();
-
-        if ( address.sourceType() != RifDataSourceForRftPlt::SUMMARY_RFT )
+        RifDataSourceForRftPlt address      = curveDefToAdd.address();
+        RifDataSourceForRftPlt colorAddress = address;
+        if ( address.sourceType() == RifDataSourceForRftPlt::SUMMARY_RFT )
         {
-            if ( !m_dataSourceColors.count( curveDefToAdd.address() ) )
-            {
-                colorTableIndex                             = colorTableIndex % colorTable.size();
-                m_dataSourceColors[curveDefToAdd.address()] = colorTable[colorTableIndex];
-            }
+            colorAddress = RifDataSourceForRftPlt( RifDataSourceForRftPlt::ENSEMBLE_RFT, address.ensemble() );
         }
 
-        if ( address.sourceType() == RifDataSourceForRftPlt::ENSEMBLE_RFT && address.ensemble() )
+        if ( !m_dataSourceColors.count( colorAddress ) )
         {
-            RimWellRftEnsembleCurveSet* curveSet = findEnsembleCurveSet( address.ensemble() );
-            CAF_ASSERT( curveSet );
-            curveSet->initializeLegend();
-
-            if ( curveSet->legendConfig()->showLegend() && curveSet->colorMode() == ColorMode::BY_ENSEMBLE_PARAM &&
-                 !curveSet->currentEnsembleParameter().isEmpty() )
-            {
-                if ( !m_ensembleLegendFrames[curveSet] )
-                {
-                    m_ensembleLegendFrames[curveSet] = new RiuCvfOverlayItemWidget( track->viewer(),
-                                                                                    track->viewer()->canvas() );
-                }
-                m_ensembleLegendFrames[curveSet]->updateFromOverlayItem( curveSet->legendConfig()->titledOverlayFrame() );
-                track->viewer()->addOverlayFrame( m_ensembleLegendFrames[curveSet] );
-            }
+            colorTableIndex                  = colorTableIndex % colorTable.size();
+            m_dataSourceColors[colorAddress] = colorTable[colorTableIndex];
         }
 
         if ( address.sourceType() != RifDataSourceForRftPlt::ENSEMBLE_RFT )
@@ -1301,7 +1316,7 @@ std::vector<RimSummaryCaseCollection*> RimWellRftPlot::selectedEnsembles() const
     std::vector<RimSummaryCaseCollection*> ensembleSets;
     for ( const RifDataSourceForRftPlt& dataSource : m_selectedSources() )
     {
-        if ( dataSource.sourceType() == RifDataSourceForRftPlt::ENSEMBLE_RFT )
+        if ( dataSource.ensemble() != nullptr )
         {
             ensembleSets.push_back( dataSource.ensemble() );
         }