nuclear_chart.C

// -----------------------------------------------------------------------------
// Nuclear chart visualizations with drip-line envelopes and stable nuclei overlays.
// -----------------------------------------------------------------------------
 
#include "tkmanager.h"
#include "tkn_root_macros.h"
 
#include <TArrow.h>
#include <TCanvas.h>
#include <TColor.h>
#include <TGraph.h>
#include <TGraphAsymmErrors.h>
#include <TH2.h>
#include <TImage.h>
#include <TLine.h>
#include <TLatex.h>
#include <TPad.h>
#include <TPaletteAxis.h>
#include <TSpline.h>
#include <TStyle.h>
#include <TPolyLine.h>
#include <TBox.h>
 
#include <cmath>
#include <algorithm>
#include <map>
#include <numeric>   // std::iota
 
using namespace tkn;
 
// ------------------------- Config / Types ------------------------------------
enum class EvenEven { No, Yes };
 
struct ChartWindow { double xmin, xmax, ymin, ymax; };
 
struct MagicLayout {
    // “Magic numbers” and label anchor ranges
    static constexpr int    Nm[7]   = {2  , 8  , 20 , 28 , 50 , 82  , 126 };
    static constexpr double NmMin[7]= {0. , 1. , 12., 18., 46., 93. , 98. };
    static constexpr double NmMax[7]= {10., 22., 40., 52., 90., 140., 126.*3};
    static constexpr double ZmMin[7]= {0.2, 0.2, 6. , 10., 26., 44. , 74.};
    static constexpr double ZmMax[7]= {9. , 17., 30., 35., 53., 74. , 95.};
};
 
// ---------------------------- Helpers ----------------------------------------
namespace {

TH2D* MakeHist(const char* name, int NMax, int ZMax, EvenEven mode) {
    if (mode == EvenEven::Yes) {
        // step=2 binning: edges at -1,1,3,... (centers on even integers)
        return new TH2D(name,";N;Z",
                        (NMax+2)/2, -1, NMax+1,
                        (ZMax+2)/2, -1, ZMax+1);
    } else {
        // unit binning centered on integers
        return new TH2D(name,";N;Z",
                        NMax+1, -0.5, NMax+0.5,
                        ZMax+1, -0.5, ZMax+0.5);
    }
}

void HideEmptyBins(TH2* h, double sentinel=-1e8) {
    for (int ix=1; ix<=h->GetNbinsX(); ++ix)
        for (int iy=1; iy<=h->GetNbinsY(); ++iy)
            if (h->GetBinContent(ix,iy)==0.0) h->SetBinContent(ix,iy, sentinel);
}

void ApplyCleanStyle(TH2* h, const ChartWindow& w, double zmin, double zmax) {
    gStyle->SetOptStat(0);
    gStyle->SetNumberContours(255);
    gStyle->SetPalette(kRainBow);
 
    h->SetMinimum(zmin);
    h->SetMaximum(zmax);
 
    h->GetXaxis()->SetRangeUser(w.xmin, w.xmax);
    h->GetYaxis()->SetRangeUser(w.ymin, w.ymax);
 
    h->SetTitle("");
    h->GetXaxis()->SetLabelSize(0);
    h->GetXaxis()->SetTitleSize(0);
    h->GetXaxis()->SetTickLength(0);
    h->GetXaxis()->SetAxisColor(0);
    h->GetYaxis()->SetLabelSize(0);
    h->GetYaxis()->SetTitleSize(0);
    h->GetYaxis()->SetTickLength(0);
    h->GetYaxis()->SetAxisColor(0);
    gPad->SetFrameLineColor(0);
}

void PlacePalette(TH2* h, double x1=0.92,double x2=0.96,double y1=0.12,double y2=0.90)
{
    gPad->Update();
 
    // White masks in pad coordinates (frame) to hide any drawings beyond axes.
    // Right of Uxmax
    auto mask = new TBox(gPad->GetUxmax(), gPad->GetY1(), gPad->GetX2(), gPad->GetY2());
    mask->SetFillColor(kWhite); mask->SetLineColor(kWhite); mask->Draw("same");
 
    // Below Uymin
    mask = new TBox(gPad->GetX1(), gPad->GetY1(), gPad->GetX2(), gPad->GetUymin());
    mask->SetFillColor(kWhite); mask->SetLineColor(kWhite); mask->Draw("same");
 
    // Left of Uxmin
    mask = new TBox(gPad->GetX1(), gPad->GetY1(), gPad->GetUxmin(), gPad->GetY2());
    mask->SetFillColor(kWhite); mask->SetLineColor(kWhite); mask->Draw("same");
 
    // Above Uymax
    mask = new TBox(gPad->GetX1(), gPad->GetUymax(), gPad->GetX2(), gPad->GetY2());
    mask->SetFillColor(kWhite); mask->SetLineColor(kWhite); mask->Draw("same");
 
    // Reposition palette and bring it to front
    if (auto pal=(TPaletteAxis*)h->GetListOfFunctions()->FindObject("palette")) {
        pal->SetX1NDC(x1); pal->SetX2NDC(x2);
        pal->SetY1NDC(y1); pal->SetY2NDC(y2);
        pal->SetLineColor(0);
        pal->SetLabelSize(0.035);
        pal->SetTitleSize(0.040);
        pal->SetTickLength(0.02);
        pal->Draw();
    }
    gPad->Modified(); gPad->Update();
}

void DrawMagicGridAndLabels(EvenEven mode,
                            int lineColor=kGray+2, int lineStyle=1, float lineWidth=2.f,
                            int textColor=kGray+3, double textSize=0.035)
{
    if (!gPad) return;
    gPad->Modified(); gPad->Update();
 
    const double nxmin = gPad->GetUxmin();
    const double nxmax = gPad->GetUxmax();
    const double nymin = gPad->GetUymin();
    const double nymax = gPad->GetUymax();
 
    const double offset = (mode==EvenEven::Yes ? 1.0 : 0.5);
    const double dx = (mode==EvenEven::Yes ? 0.045 : 0.03) * (nxmax - nxmin);
    const double dy = (mode==EvenEven::Yes ? 0.045 : 0.03) * (nymax - nymin);
 
    TLatex lab; lab.SetTextFont(42); lab.SetTextColor(textColor); lab.SetTextSize(textSize);
 
    // Vertical pairs at N = m ± offset
    for (int i=0; i<7; ++i) {
        const int    N  = MagicLayout::Nm[i];
        const double z1 = std::max(nymin, MagicLayout::ZmMin[i]);
        const double z2 = std::min(nymax, MagicLayout::ZmMax[i]);
 
        if (N >= nxmin + offset && N <= nxmax - offset && z1 < nymax) {
            for (double s : {-offset, +offset}) {
                auto L = new TLine(N + s, z1, N + s, z2);
                L->SetLineColor(lineColor); L->SetLineStyle(lineStyle); L->SetLineWidth(lineWidth);
                L->Draw();
            }
            lab.SetTextAlign(22);
            lab.DrawLatex(N, z1 - dy, Form("%d", N));
        }
    }
 
    // Horizontal pairs at Z = m ± offset
    for (int i=0; i<7; ++i) {
        const int    Z  = MagicLayout::Nm[i];
        const double n1 = std::max(nxmin, MagicLayout::NmMin[i]);
        const double n2 = std::min(nxmax, MagicLayout::NmMax[i]);
 
        if (Z >= nymin + offset && Z <= nymax - offset && n1 < nxmax) {
            for (double s : {-offset, +offset}) {
                auto L = new TLine(n1, Z + s, n2, Z + s);
                L->SetLineColor(lineColor); L->SetLineStyle(lineStyle); L->SetLineWidth(lineWidth);
                L->Draw("same");
            }
            lab.SetTextAlign(22);
            lab.DrawLatex(n1 - dx, Z, Form("%d", Z));
        }
    }
}

void DrawAxisHints() {
    TArrow arrow; arrow.SetLineColorAlpha(kGray+2, 0.9); arrow.SetLineWidth(2);
 
    auto arr = arrow.DrawArrow(0.55, 0.045, 0.65, 0.045, 0.012, "|>"); arr->SetNDC();
    TLatex tx; tx.SetTextFont(42); tx.SetTextSize(0.04); tx.SetTextColor(kGray+2); tx.SetTextAlign(22);
    tx.DrawLatexNDC(0.67, 0.05, "N");
 
    arr = arrow.DrawArrow(0.022, 0.5, 0.022, 0.65, 0.012, "|>"); arr->SetNDC();
    TLatex tz; tz.SetTextFont(42); tz.SetTextSize(0.04); tz.SetTextColor(kGray+2); tz.SetTextAlign(22);
    tz.SetTextAngle(90);
    tz.DrawLatexNDC(0.02, 0.67, "Z");
    tz.SetTextAngle(0);
}

void DrawTitleTopLeft(const char* txt) {
    const double xL=gPad->GetLeftMargin(), yT=1.0-gPad->GetTopMargin();
    TLatex t; t.SetNDC(true); t.SetTextFont(42); t.SetTextColor(kGray+3); t.SetTextSize(0.045); t.SetTextAlign(13);
    t.DrawLatex(xL+0.010, yT-0.010, txt);
}

static void DrawSmoothBandX(const std::vector<double>& x,   // Nmid
                            const std::vector<double>& y,   // Z
                            const std::vector<double>& ex,  // half-width in N
                            Color_t fillColor, float alpha = 0.25,
                            Color_t lineColor = kBlack, int lineWidth = 2,
                            bool smooth = true, int nsamples = -1)
{
    const int n = (int)x.size();
    if (n < 2) return;
 
    // Sort by increasing y (Z) so we build monotone x=f(y)
    std::vector<int> idx(n); std::iota(idx.begin(), idx.end(), 0);
    std::sort(idx.begin(), idx.end(), [&](int a, int b){ return y[a] < y[b]; });
 
    std::vector<double> yS(n), xHi(n), xLo(n);
    for (int i=0;i<n;++i) {
        const int k = idx[i];
        yS[i]  = y[k];
        xHi[i] = x[k] + ex[k];
        xLo[i] = x[k] - ex[k];
    }
 
    std::vector<double> xp, yp; // polygon
 
    if (!smooth) {
        // ------------------------ STEP (STAIR) ENVELOPE ------------------------
        // Build "stair" polylines for upper and lower borders.
        // Algorithm (for x=f(y)):
        //   points: (x0,y0), then for each i=0..n-2:
        //     (x_i,   y_{i+1})  // vertical segment at x_i
        //     (x_{i+1},y_{i+1}) // horizontal segment to new x
        auto build_stairs = [](const std::vector<double>& yv,
                               const std::vector<double>& xv,
                               std::vector<double>& outX,
                               std::vector<double>& outY)
        {
            const int m = (int)xv.size();
            outX.clear(); outY.clear();
            outX.reserve(2*m); outY.reserve(2*m);
            outX.push_back(xv[0]); outY.push_back(yv[0]);
            for (int i=0;i<m-1;++i) {
                outX.push_back(xv[i]);     outY.push_back(yv[i+1]); // vertical
                outX.push_back(xv[i+1]);   outY.push_back(yv[i+1]); // horizontal
            }
        };
 
        std::vector<double> xU, yU, xL, yL;
        build_stairs(yS, xHi, xU, yU);
        build_stairs(yS, xLo, xL, yL);
 
        // Build filled polygon: upper (forward) + lower (reverse)
        xp.reserve(xU.size() + xL.size());
        yp.reserve(yU.size() + yL.size());
        for (size_t i=0;i<xU.size();++i) { xp.push_back(xU[i]); yp.push_back(yU[i]); }
        for (int i=(int)xL.size()-1;i>=0;--i) { xp.push_back(xL[i]); yp.push_back(yL[i]); }
 
        // Draw fill
        auto* poly = new TPolyLine((int)xp.size(), xp.data(), yp.data());
        poly->SetFillColorAlpha(fillColor, alpha);
        poly->SetLineColor(lineColor);
        poly->SetLineWidth(lineWidth);
        poly->Draw("f");
 
        // Draw borders as step polylines
        auto* grHi = new TGraph((int)xU.size());
        auto* grLo = new TGraph((int)xL.size());
        for (int i=0;i<(int)xU.size();++i) grHi->SetPoint(i, xU[i], yU[i]);
        for (int i=0;i<(int)xL.size();++i) grLo->SetPoint(i, xL[i], yL[i]);
        grHi->SetLineColor(lineColor); grHi->SetLineWidth(lineWidth); grHi->Draw("L");
        grLo->SetLineColor(lineColor); grLo->SetLineWidth(lineWidth); grLo->Draw("L");
    } else {
        // --------------------------- SMOOTH ENVELOPE ---------------------------
        TSpline3 sHi("sHi", yS.data(), xHi.data(), n);
        TSpline3 sLo("sLo", yS.data(), xLo.data(), n);
 
        const int m = (nsamples>1) ? nsamples : std::max(100, 4*n);
        const double yMin = yS.front(), yMax = yS.back();
 
        xp.reserve(2*m); yp.reserve(2*m);
        for (int i=0;i<m;++i) {
            const double yy = yMin + (yMax - yMin) * (double(i)/double(m-1));
            xp.push_back(sHi.Eval(yy));
            yp.push_back(yy);
        }
        for (int i=m-1;i>=0;--i) {
            const double yy = yMin + (yMax - yMin) * (double(i)/double(m-1));
            xp.push_back(sLo.Eval(yy));
            yp.push_back(yy);
        }
 
        auto* poly = new TPolyLine((int)xp.size(), xp.data(), yp.data());
        poly->SetFillColorAlpha(fillColor, alpha);
        poly->SetLineColor(lineColor);
        poly->SetLineWidth(lineWidth);
        poly->Draw("f");
 
        auto* grHi = new TGraph(m);
        auto* grLo = new TGraph(m);
        for (int i=0;i<m;++i) {
            grHi->SetPoint(i, xp[i],             yp[i]);
            grLo->SetPoint(i, xp[2*m-1-i],       yp[2*m-1-i]);
        }
        grHi->SetLineColor(lineColor); grHi->SetLineWidth(lineWidth); grHi->Draw("L");
        grLo->SetLineColor(lineColor); grLo->SetLineWidth(lineWidth); grLo->Draw("L");
    }
}

static void BuildEnvelopeCenterHalfwidth(
    const std::vector<tkn_drip_point>& pts,
    double minHalfBin,                           // 0.5 for unit bins; 1.0 for even-even
    std::vector<double>& xNmid,                  // OUT: Nmid
    std::vector<double>& yZ,                     // OUT: Z
    std::vector<double>& exHalfWidth             // OUT: half-width on N
    ){
    std::map<int, std::pair<int,int>> z2mm;     // Z -> (Nmin,Nmax)
    for (const auto& p : pts) {
        auto& mm = z2mm[p.Z];
        if (mm.first==0 && mm.second==0) mm = {p.N,p.N};
        mm.first  = (mm.first==0  ? p.N : std::min(mm.first,  p.N));
        mm.second = (mm.second==0 ? p.N : std::max(mm.second, p.N));
    }
    xNmid.clear(); yZ.clear(); exHalfWidth.clear();
    xNmid.reserve(z2mm.size()); yZ.reserve(z2mm.size()); exHalfWidth.reserve(z2mm.size());
 
    for (const auto& [Z, mm] : z2mm) {
        const double Nmid = 0.5*(mm.first + mm.second);
        double halfW = 0.5*(mm.second - mm.first);
        if (halfW < minHalfBin) halfW = minHalfBin; // keep visible for single N
        xNmid.push_back(Nmid);
        yZ.push_back((double)Z);
        exHalfWidth.push_back(halfW);
    }
}

static TGraph* MakeDripLineFromPoints(const std::vector<tkn_drip_point>& pts,
                                      Color_t color, Style_t lineStyle = 1, Width_t lineWidth = 3,
                                      Style_t markerStyle = 20, Size_t markerSize = 0)
{
    if (pts.empty()) return nullptr;
    std::vector<tkn_drip_point> v = pts;
    std::sort(v.begin(), v.end(), [](auto& a, auto& b){ return (a.Z<b.Z)||(a.Z==b.Z && a.N<b.N); });
 
    std::vector<double> xs, ys; xs.reserve(v.size()); ys.reserve(v.size());
    for (const auto& p : v) { xs.push_back(p.N); ys.push_back(p.Z); }
 
    auto* gr = new TGraph((int)xs.size(), xs.data(), ys.data());
    gr->SetLineColor(color); gr->SetLineStyle(lineStyle); gr->SetLineWidth(lineWidth);
    gr->SetMarkerStyle(markerStyle); if (markerSize>0) gr->SetMarkerSize(markerSize);
    gr->SetMarkerColor(color);
    return gr;
}

void DrawDripLinesIfAvailable(EvenEven mode)
{
    const double minHalfBin = (mode==EvenEven::No) ? 0.5 : 1.0;
 
    // All values (for envelopes)
    auto s2n_all = gmanager->get_drip_line("S2n", false);
    auto s2p_all = gmanager->get_drip_line("S2p", false);
 
    // Build & draw smooth envelopes
    std::vector<double> xNmid, yZ, exHalf;
    if (!s2n_all.empty()) {
        BuildEnvelopeCenterHalfwidth(s2n_all, minHalfBin, xNmid, yZ, exHalf);
        DrawSmoothBandX(xNmid, yZ, exHalf,
                        /*fillColor=*/kAzure+1, 0.25,
                        /*lineColor=*/kAzure+2, 2,
                        /*smooth=*/true);
    }
    if (!s2p_all.empty()) {
        BuildEnvelopeCenterHalfwidth(s2p_all, minHalfBin, xNmid, yZ, exHalf);
        DrawSmoothBandX(xNmid, yZ, exHalf,
                        /*fillColor=*/kRed-7, 0.25,
                        /*lineColor=*/kRed+1, 2,
                        /*smooth=*/true);
    }
 
    // Optional center lines (min |Q|) if desired:
    // auto s2n_min = gmanager->get_drip_line("S2n", true);
    // if (auto* gS2n = MakeDripLineFromPoints(s2n_min, kAzure+2, 1, (mode==EvenEven::No)?4:3)) gS2n->Draw("L");
    // auto s2p_min = gmanager->get_drip_line("S2p", true);
    // if (auto* gS2p = MakeDripLineFromPoints(s2p_min, kRed+1,   1, (mode==EvenEven::No)?4:3)) gS2p->Draw("L");
 
    // Simple legend
    TLatex lab; lab.SetTextFont(42); lab.SetTextSize(0.04); lab.SetTextAlign(13);
    lab.SetTextColor(kAzure+2);
    lab.DrawLatexNDC(gPad->GetLeftMargin()+0.010, 0.90, "S_{2n} drip line");
    lab.SetTextColor(kRed+1);
    lab.DrawLatexNDC(gPad->GetLeftMargin()+0.010, 0.86, "S_{2p} drip line");
}

static void DrawStableNucleiBoxes(EvenEven mode,
                                  Color_t color = kGray+2,
                                  double alpha = 1.0,
                                  Width_t lineWidth = 2,
                                  double inset = 0.08,
                                  bool fill = false)
{
    const double halfBin = (mode == EvenEven::Yes) ? 1.0 : 0.5;
 
    glog.set_warnings(false);
    for (const auto& nuc : gmanager->get_nuclei()) {
        const int Z = nuc->get_z();
        const int N = nuc->get_n();
 
        if (mode == EvenEven::Yes && ((N % 2) || (Z % 2))) continue;
        if (!nuc->is_stable()) continue;
 
        const double x1 = N - halfBin + inset;
        const double x2 = N + halfBin - inset;
        const double y1 = Z - halfBin + inset;
        const double y2 = Z + halfBin - inset;
 
        auto* box = new TBox(x1, y1, x2, y2);
        box->SetLineColorAlpha(color, alpha);
        box->SetLineWidth(lineWidth);
 
        if (fill) {
            box->SetFillStyle(1001);
            box->SetFillColorAlpha(color, alpha);
        } else {
            box->SetFillStyle(0); // hollow
        }
        box->Draw("same");
    }
}
 
} // namespace
 
// ============================= PLOT 1: B(E2)/A ================================
void plot_nuclear_chart_deformation()
{
    EvenEven mode = EvenEven::Yes;
    const ChartWindow W{18,130, 18,85};
    const int NMax=180, ZMax=120;
 
    TH2D* h = MakeHist("NucChart_BE2", NMax, ZMax, mode);
 
    // --- Fill (even-even filter only if mode == Yes)
    glog.set_warnings(false);
    for (const auto &nuc : gmanager->get_nuclei([](auto nuc) {
             return (nuc->get_z()%2==0 && nuc->get_n()%2==0);
         }))
    {
        auto lev_scheme = nuc->get_level_scheme();
        auto gamma = lev_scheme->get_decay<tkgammadecay>("2+1->0+1",false);
        if(!gamma) continue;
        auto BE2W = gamma->get_trans_prob(true,2,true);
        if(BE2W<=0||std::isnan(BE2W)) continue;
        h->Fill(nuc->get_n(), nuc->get_z(), BE2W / ((float)nuc->get_a()));
    }
 
    HideEmptyBins(h);
 
    // Canvas + style + draw
    auto* c = new TCanvas("c_be2","",1000,600);
    c->SetMargin(0.06, 0.09, 0.08, 0.04);
 
    ApplyCleanStyle(h, W, /*zmin*/0.0, /*zmax*/2.0);
    h->Draw("COLZ");
 
    DrawDripLinesIfAvailable(EvenEven::Yes);
    DrawStableNucleiBoxes(/*mode=*/EvenEven::Yes, /*color=*/kGray+1, /*alpha=*/1.0,
                          /*lineWidth=*/1, /*inset=*/0.20);
    PlacePalette(h, 0.92, 0.96, c->GetBottomMargin(), 1-c->GetTopMargin());
    DrawMagicGridAndLabels(mode);
 
    DrawAxisHints();
    DrawTitleTopLeft("B(E2)/A (W.u.)");
 
    tkn::DrawLogoAndCredit(0.8,0.08,0.1,0.027,"Generated with","NDC");
 
    gPad->Modified(); gPad->Update();
 
    c->SaveAs("nuclear_chart_deformation.png");
}
 
// ============ PLOT 2: Double differential of mass excess Δm ===================
void plot_nuclear_chart_double_mass_diff()
{
    EvenEven mode = EvenEven::No;
    const ChartWindow W{18,130, 18,85};
    const int NMax=180, ZMax=120;
 
    TH2D* h = MakeHist("NucChart_DD", NMax, ZMax, mode);
 
    glog.set_warnings(false);
    const int dn=2, dz=2;
    for (const auto &nuc : gmanager->get_nuclei()) {
 
        tknucleus n1(nuc->get_z()-dz, nuc->get_a()-dn-dz);
        tknucleus n2(nuc->get_z()+dz, nuc->get_a()+dn+dz);
        if (!n1.is_known() || !n2.is_known()) continue;
 
        const double me  = nuc->get_mass_excess(tkunit_manager::MeV);
        const double me1 = n1.get_mass_excess(tkunit_manager::MeV);
        const double me2 = n2.get_mass_excess(tkunit_manager::MeV);
        if (std::isnan(me) || std::isnan(me1) || std::isnan(me2)) continue;
 
        const double dd = (me1 - 2.0*me + me2) / (2.0*std::sqrt(dn*dn + dz*dz));
        h->SetBinContent(h->FindBin(nuc->get_n(), nuc->get_z()), dd);
    }
 
    HideEmptyBins(h);
 
    auto* c = new TCanvas("c_dd","",1000,600);
    c->SetMargin(0.06, 0.09, 0.08, 0.04);
 
    ApplyCleanStyle(h, W, /*zmin*/-1.0, /*zmax*/2.5);
    h->Draw("COLZ");
 
    DrawDripLinesIfAvailable(EvenEven::Yes);
    DrawStableNucleiBoxes(/*mode=*/EvenEven::No, /*color=*/kBlack, /*alpha=*/0.6,
                          /*lineWidth=*/1, /*inset=*/0.0, /*fill=*/true);
    PlacePalette(h, 0.92, 0.96, c->GetBottomMargin(), 1-c->GetTopMargin());
    DrawMagicGridAndLabels(mode);
 
    DrawAxisHints();
    DrawTitleTopLeft("Double differential of the mass excess");
 
    tkn::DrawLogoAndCredit(0.8,0.08,0.1,0.027,"Generated with","NDC");
 
    gPad->Modified(); gPad->Update();
 
    c->SaveAs("nuclear_chart_double_mass_diff.png");
}