#include "TChain.h"
#include "TH2D.h"
#include "TCanvas.h"
#include "TTreeReader.h"
#include "TTreeReaderValue.h"
#include <array>
#include <string>
#include <algorithm>
#include <iostream>

namespace meta
{
    struct value
    {
        value(int id, int bin, std::string name)
            : id(id)
            , bin(bin)
            , name(name)
        {}

        int id;
        int bin;
        std::string name;
    };
}

void reach()
{
    // load data 
    TChain e("e");
    e.AddFile("/Users/rwkelley/Development/dpm/research/optimization/data/pagv.root");

    // register branches (i.e. columns)
    TTreeReader reader(&e);
    //TTreeReaderValue<long long> entity_id           (reader, "entity_id"              );
    TTreeReaderValue<long long> network_id          (reader, "network_id"             );
    TTreeReaderValue<long long> daypart_id          (reader, "daypart_id"             );
    TTreeReaderValue<long long> program_airing_group(reader, "program_airing_group_id");
    TTreeReaderValue<double   > avg_fraction_viewed (reader, "avg_fraction_viewed"    );
    TTreeReaderValue<double   > weight              (reader, "weight"                 );

    // meta-data 
    std::array<meta::value const, 5> const networks =
    {
        meta::value{1  , 1, "ABC"  },
        meta::value{9  , 2, "BRAVO"},
        meta::value{24 , 3, "ESPN" },
        meta::value{65 , 4, "NBC"  },
        meta::value{107, 5, "USA"  }
    };


    std::array<meta::value const, 6> const dayparts =
    {
        meta::value{1, 1, "Morning"      },
        meta::value{2, 2, "Daytime"      },
        meta::value{3, 3, "Early-Fringe" },
        meta::value{4, 4, "Prime"        },
        meta::value{5, 5, "Late-Fringe"  },
        meta::value{6, 6, "Overnight"    } 
    };

    // hold the answer in a histogram
    TH2D * const h_impressions = new TH2D("h_impressions", "Impressions;daypart;network", dayparts.size(), 0, dayparts.size(), networks.size(), 0, networks.size());

    // style histogram axes
    h_impressions->SetStats(0);
    for (auto const &daypart : dayparts)
    {
        h_impressions->GetXaxis()->SetBinLabel(daypart.bin, daypart.name.c_str());
    }
    for (auto const &network : networks)
    {
        h_impressions->GetYaxis()->SetBinLabel(network.bin, network.name.c_str());
    }

    // loop
    //reader.SetEntriesRange(0, 100);
    while (reader.Next())
    {
        // local variables
        int const &n_id = *network_id;
        int const &d_id = *daypart_id;

        // filter for network
        auto network_itr = std::find_if(networks.begin(), networks.end(), [n_id](meta::value p){return p.id==n_id;});
        if (network_itr == networks.end())
            continue;

        // filter for daypart
        auto daypart_itr = std::find_if(dayparts.begin(), dayparts.end(), [d_id](meta::value p){return p.id==d_id;});
        if (daypart_itr == dayparts.end())
            continue;

        // fill the hist
        meta::value const& network = *network_itr;
        meta::value const& daypart = *daypart_itr;
        double const f     = *avg_fraction_viewed;
        double const w     = *weight;
        h_impressions->Fill(daypart.name.c_str(), network.name.c_str(), w*f);
    }

    // display
    TCanvas * const c = new TCanvas;
    h_impressions->Draw("colz");
}