"""
Figures for:
"Strength reduction and strength-scaling relations of M40 concrete with
untreated bamboo chips as partial coarse-aggregate replacement"

  Fig1  strength retained vs bamboo content
  Fig2  flexural and splitting tensile strength vs compressive strength (log-log)
  Fig3  compressive strength retained compared with published data

Each figure is saved to ./figures as EPS and PDF (vector, fonts embedded)
and as TIFF (LZW) and PNG at DPI resolution. Adjust the SETTINGS block to a
journal's figure requirements (width ~84-90 mm for one column of a two-column
page, ~174-190 mm for full page width).

Requires: Python 3.8+, numpy, scipy, matplotlib
Run:      python make_figures.py
"""
import os

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from scipy import stats

from bamboo_data import LEVELS, COMPRESSIVE, SPLITTING, FLEXURAL, LITERATURE

# ----------------------------------------------------------------- SETTINGS
FIG_W_MM = 174          # figure width (mm)
FONT_STYLE = "sans"     # "sans" (Arial/Helvetica) or "serif" (Times)
DPI = 600               # resolution of TIFF/PNG files
FILE_PREFIX = "Fig"     # files are named Fig1, Fig2, Fig3

# -------------------------------------------------------------------- STYLE
MM = 1 / 25.4
NARROW = FIG_W_MM < 120              # narrow figures: smaller text, legend below
FS = 7 if NARROW else 9              # base font size (pt)
FIGSIZE = (FIG_W_MM * MM, (0.80 if NARROW else 0.66) * FIG_W_MM * MM)
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "figures")

mpl.rcParams.update({
    "font.family": "sans-serif" if FONT_STYLE == "sans" else "serif",
    "font.sans-serif": ["Arial", "Helvetica", "Liberation Sans", "DejaVu Sans"],
    "font.serif": ["Times New Roman", "Liberation Serif", "DejaVu Serif"],
    "mathtext.fontset": "stixsans" if FONT_STYLE == "sans" else "stix",
    "font.size": FS,
    "axes.labelsize": FS,
    "xtick.labelsize": FS - 0.5,
    "ytick.labelsize": FS - 0.5,
    "legend.fontsize": FS - 1,
    "legend.frameon": False,
    "axes.linewidth": 0.6,
    "xtick.major.width": 0.6,
    "ytick.major.width": 0.6,
    "xtick.direction": "in",
    "ytick.direction": "in",
    "xtick.top": True,
    "ytick.right": True,
    "savefig.dpi": DPI,
    "savefig.bbox": "tight",
    "savefig.pad_inches": 0.02,
    "pdf.fonttype": 42,
    "ps.fonttype": 42,
})

# Colour-blind-safe palette; each series also has its own marker and line
# style so the figures remain readable in greyscale.
BLACK, BLUE, VERMILION = "#000000", "#0072B2", "#D55E00"
GREEN, ORANGE, GREY = "#009E73", "#E69F00", "#555555"
FC = r"f^{\prime}_{\!c}"
DOTTED, DASHED = (0, (1.2, 1.6)), (0, (5, 2.5))


def mean_sd(data):
    return (np.array([np.mean(data[k]) for k in LEVELS]),
            np.array([np.std(data[k], ddof=1) for k in LEVELS]))


p = np.array(LEVELS)
fc, fc_sd = mean_sd(COMPRESSIVE)
fsp, fsp_sd = mean_sd(SPLITTING)
fr, fr_sd = mean_sd(FLEXURAL)


def save(fig, number):
    name = os.path.join(OUT, f"{FILE_PREFIX}{number}")
    for ext in ("eps", "pdf", "png"):
        fig.savefig(f"{name}.{ext}")
    fig.savefig(f"{name}.tiff", pil_kwargs={"compression": "tiff_lzw"})
    plt.close(fig)


def add_legend(ax, handles=None, labels=None, loc="upper left"):
    if handles is None:
        handles, labels = ax.get_legend_handles_labels()
    if NARROW:
        ax.legend(handles, labels, loc="upper center", bbox_to_anchor=(0.5, -0.2),
                  ncol=2, handlelength=2.2, columnspacing=1.0)
    else:
        ax.legend(handles, labels, loc=loc, handlelength=2.5)


def fig1():
    fig, ax = plt.subplots(figsize=FIGSIZE)
    series = [(rf"Compressive, ${FC}$", fc, fc_sd, BLACK, "o", "-", BLACK),
              ("Flexural, $f_r$", fr, fr_sd, BLUE, "s", "--", "white"),
              ("Splitting tensile, $f_{sp}$", fsp, fsp_sd, VERMILION, "^", ":", "white")]
    for label, m, sd, c, marker, ls, mfc in series:
        retained = 100 * m / m[0]
        ax.errorbar(p, retained, yerr=100 * sd / m[0], color=c, ls=ls, marker=marker,
                    mfc=mfc, mec=c, mew=0.9, ms=5.5, capsize=2.5, elinewidth=0.7,
                    label=label, zorder=3)
        ax.annotate(f"{retained[-1]:.0f}%", xy=(p[-1], retained[-1]), xytext=(7, 0),
                    textcoords="offset points", va="center", fontsize=FS - 1, color=c)
    ax.set(xlim=(-0.8, 18), ylim=(0, 105), xticks=p, yticks=np.arange(0, 101, 20),
           xlabel="Bamboo chips (% of coarse aggregate by volume)",
           ylabel="Strength retained (% of control)")
    add_legend(ax, loc="lower left")
    save(fig, 1)


def fig2():
    fig, ax = plt.subplots(figsize=FIGSIZE)
    xx = np.linspace(13.5, 47, 200)
    for m, sd, label, marker, c, lx, dy, va, ha in [
            (fr, fr_sd, "Flexural, $f_r$", "s", BLUE, 26, 1.08, "bottom", "right"),
            (fsp, fsp_sd, "Splitting tensile, $f_{sp}$", "^", VERMILION, 31, 0.87, "top", "left")]:
        fit = stats.linregress(np.log(fc), np.log(m))
        a, b = np.exp(fit.intercept), fit.slope
        k = np.mean(m / np.sqrt(fc))
        ax.errorbar(fc, m, xerr=fc_sd, yerr=sd, ls="none", marker=marker, mfc="white",
                    mec=c, ecolor=c, mew=0.9, ms=5.5, elinewidth=0.6, capsize=2,
                    label=label, zorder=4)
        ax.plot(xx, a * xx ** b, color=c, lw=1.1, zorder=3)
        ax.plot(xx, k * np.sqrt(xx), color=GREY, lw=0.9, ls=DOTTED, zorder=2)
        symbol = "f_r" if c == BLUE else "f_{sp}"
        ax.text(lx, a * lx ** b * dy,
                rf"${symbol} = {a:.3f}\,({FC})^{{{b:.2f}}}$" + "\n"
                + rf"$R^2 = {fit.rvalue ** 2:.3f}$",
                color=c, fontsize=FS - 1, va=va, ha=ha)
    ax.plot(xx, 0.62 * np.sqrt(xx), color=GREY, lw=0.9, ls=DASHED, zorder=2)

    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set(xlim=(13.5, 47), ylim=(1.4, 6), xticks=[15, 20, 25, 30, 40],
           yticks=[1.5, 2, 3, 4, 5, 6],
           xlabel=rf"Compressive strength, ${FC}$ (MPa)",
           ylabel="Flexural or splitting tensile strength (MPa)")
    for axis in (ax.xaxis, ax.yaxis):
        axis.set_major_formatter(mpl.ticker.FuncFormatter(lambda v, _: f"{v:g}"))
        axis.set_minor_formatter(mpl.ticker.NullFormatter())

    handles, labels = ax.get_legend_handles_labels()
    handles += [Line2D([], [], color=BLACK, lw=1.1),
                Line2D([], [], color=GREY, lw=0.9, ls=DOTTED),
                Line2D([], [], color=GREY, lw=0.9, ls=DASHED)]
    labels += [rf"Power-law fit, $a\,({FC})^{{b}}$",
               rf"$k\sqrt{{{FC}}}$ (exponent 0.5)",
               rf"ACI 318: $0.62\sqrt{{{FC}}}$"]
    add_legend(ax, handles, labels, loc="upper left")
    save(fig, 2)


def fig3():
    fig, ax = plt.subplots(figsize=FIGSIZE)
    styles = [(GREEN, "^", "--"), (ORANGE, "s", "--"), (VERMILION, "v", "--"), (BLUE, "D", ":")]
    ax.plot(p, 100 * fc / fc[0], color=BLACK, marker="o", ls="-", lw=1.8, ms=5.5,
            label="This study (M40, untreated chips)", zorder=4)
    for (label, (control, values)), (c, marker, ls) in zip(LITERATURE.items(), styles):
        x = [0] + list(values)
        y = [100] + [100 * v / control for v in values.values()]
        ax.plot(x, y, color=c, marker=marker, ls=ls, mfc="white", mec=c, mew=0.9,
                ms=5, lw=1.0, label=label, zorder=2)
    ax.set(xlim=(-1, 42), ylim=(0, 105), xticks=np.arange(0, 41, 5),
           yticks=np.arange(0, 101, 20),
           xlabel="Bamboo replacing coarse aggregate (%)",
           ylabel="Compressive strength retained (% of control)")
    add_legend(ax, loc="upper right")
    save(fig, 3)


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    fig1()
    fig2()
    fig3()
    print(f"Figures saved to {OUT}")
