from scipy.constants import physical_constants
from TRIOMA.tools.TriomaClass import TriomaClass
from scipy.constants import N_A
from TRIOMA.tools.Extractors.PipeSubclasses import Fluid
import matplotlib.pyplot as plt


class BreedingBlanket(TriomaClass):
    """
    Represents a breeding blanket component in a fuel cycle system.

    Attributes:
        c_in(float): Inlet concentration of tritium in the breeding blanket component.
        Q (float): Heat generated by the breeding blanket component.
        TBR (float): Tritium breeding ratio of the breeding blanket component.
        T_out (float): Outlet temperature of the breeding blanket component.
        T_in (float): Inlet temperature of the breeding blanket component.
        fluid (Fluid): Fluid used in the breeding blanket component.
    """

    def __init__(
        self,
        c_in: float = None,
        Q: float = None,
        TBR: float = None,
        T_out: float = None,
        T_in: float = None,
        fluid: Fluid = None,
        name: str = None,
        m_coolant: float = None,
    ):
        self.c_in = c_in
        self.Q = Q
        self.TBR = TBR
        self.T_out = T_out
        self.T_in = T_in
        self.fluid = fluid
        self.name = name
        self.m_coolant = m_coolant

    def plot_component(self):
        fig, ax2 = plt.subplots(1, 1, figsize=(10, 5))
        rectangle = plt.Rectangle(
            (0.2, 0.3), 0.55, 0.4, edgecolor="black", facecolor="green", alpha=0.5
        )
        ax2.add_patch(rectangle)
        # Arrow pointing to the left side of the rectangle
        ax2.arrow(
            0.5, 0.7, 0, 0.1, head_width=0.05, head_length=0.1, fc="black", ec="black"
        )
        # Arrow pointing out of the right side of the rectangle
        ax2.arrow(
            0.5, 0.1, 0.0, 0.1, head_width=0.05, head_length=0.1, fc="black", ec="black"
        )
        ax2.set_aspect("equal")
        ax2.set_xlim(0, 1)
        ax2.set_ylim(0, 1)
        if self.name is None:
            ax2.set_title("Component  ")
        else:
            ax2.set_title(self.name)

        # Add text over the arrows
        ax2.text(
            0.7,
            0.8,
            r"$T_o$=" + str(self.T_out) + " K",
            color="black",
            ha="center",
            va="center",
        )
        ax2.text(
            0.7,
            0.2,
            r"$T_i$=" + str(self.T_in) + " K",
            color="black",
            ha="center",
            va="center",
        )
        ax2.text(
            0.3,
            0.2,
            f"$c_i$={self.c_in:.4g} $mol/m^3$",
            color="black",
            ha="center",
            va="center",
        )
        ax2.text(
            0.5, 0.6, f"Q={self.Q/1E6:.3g} MW", color="black", ha="center", va="center"
        )
        ax2.text(
            0.5, 0.4, f"TBR={self.TBR:.3g}", color="black", ha="center", va="center"
        )

        ax2.text(
            0.3,
            0.8,
            rf"$c_o$={self.c_out:.4g}$mol/m^3$",
            color="black",
            ha="center",
            va="center",
        )
        ax2.axis("off")
        # Display the plot
        fig.tight_layout()
        return fig

    def get_flowrate(self):
        """
        Calculates the flow rate of the coolant in the breeding blanket component.
        """
        self.m_coolant = self.Q / ((self.T_out - self.T_in) * self.fluid.cp)
        return

    # def connect_to_component(self, component2: Union["Component", "BreedingBlanket"]):
    #     component2.update_attribute("c_in", self.c_out)
    def connect_to_component(self):
        return

    def get_cout(self, print_var: bool = False):
        """
        Calculates the outlet concentration of tritium in the breeding blanket component.

        Args:
            print_var (bool): If True, prints the intermediate variables.

        Returns:
            None
        """
        if self.m_coolant is None:
            self.get_flowrate()
        eV_to_J = physical_constants["electron volt-joule relationship"][0]
        reaction_energy = 17.6e6  # reaction energy in eV 17.6 MeV
        neutrons = self.Q / (reaction_energy * eV_to_J)

        tritium_gen = self.TBR * neutrons / N_A  ##moles
        if print_var:
            print("neu", neutrons)
            print("Trit", tritium_gen)
        if self.fluid.MS:
            self.c_out = (
                tritium_gen / 2 / (self.m_coolant / self.fluid.rho) + self.c_in
            )  ##concentration of Q2 in molten salts
        else:
            self.c_out = (
                tritium_gen / (self.m_coolant / self.fluid.rho) + self.c_in
            )  ## concentration of Q in liquid metals
