#!/usr/bin/env python3

import numpy as np
import torch
import re
from steammaterials.STEAM_materials import STEAM_materials

class MaterialProperties:
    """
    A library containing all material property selector functions.
    These functions provide density, resistivity, conductivity, and specific heat
    calculations for various materials used in coil simulations.

    STEAM_materials objects are cached to avoid repeated instantiation overhead.
    """

    def __init__(self):
        """
        Initialize the MaterialProperties class with a cache for STEAM_materials.

        The cache significantly improves performance by avoiding repeated
        instantiation of STEAM_materials objects, which was a major bottleneck
        consuming ~229 seconds (8.7% of runtime) in profiling tests.
        """
        # Cache for STEAM_materials objects
        # Key: (fit_name, max_points) -> Value: STEAM_materials instance
        self._steam_cache = {}

    def get_cached_steam_material(
        self,
        fit_name: str,
        max_points: int = None
    ) -> STEAM_materials:
        """
        Get a cached STEAM_materials object or create and cache a new one.

        Args:
            fit_name: Name of the STEAM fit function
            max_points: Optional maximum number of points for buffer pre-allocation.
                       The wrapper automatically determines array dimensions at runtime.

        Returns:
            STEAM_materials: Cached or newly created STEAM_materials instance

        Performance:
            This caching eliminates redundant object creations that were observed
            in profiling, saving ~200+ seconds of runtime.
        """
        cache_key = (fit_name, max_points)

        if cache_key not in self._steam_cache:
            # Create and cache the new STEAM_materials object
            if max_points is not None:
                self._steam_cache[cache_key] = STEAM_materials(
                    fit_name, max_points=max_points
                )
            else:
                self._steam_cache[cache_key] = STEAM_materials(fit_name)

        return self._steam_cache[cache_key]

    @staticmethod
    def get_fit_input_parameters(fit_name: str) -> list[str]:
        """
        Get the required input parameters for a given fit function.
        
        Args:
            fit_name (str): Name of the fit function (e.g., 'CFUN_HTS_JcFit_Fujikura_v1')
            
        Returns:
            list[str]: List of required input parameter names
            
        Raises:
            ValueError: If the fit function name is not recognized
        """
        # Dictionary mapping fit function names to their required input parameters
        fit_parameters = {
            # Specific heat functions - Temperature only
            'CFUN_CvCu_v1': ['Temperature'],
            'CFUN_CvAlAlloy_v1': ['Temperature'],
            'CFUN_CvAl_v1': ['Temperature'],
            'CFUN_CvHast_v1': ['Temperature'],
            'CFUN_CvAg_v1': ['Temperature'],
            'CFUN_CvSteel_v1': ['Temperature'],
            'CFUN_CvSS316_NIST_v1': ['Temperature'],
            
            # Resistivity functions
            'CFUN_rhoCu_NIST_v2': ['Temperature', 'Magnetic Field'],  # Also uses RRR internally
            'CFUN_rhoAl6061_v1': ['Temperature'],
            'CFUN_rhoAl1350_v1': ['Temperature'],
            'CFUN_rhoHast_v2': ['Temperature'],
            'CFUN_rhoAg_v1': ['Temperature', 'Magnetic Field'],  # Also uses RRR internally
            'CFUN_rhoSS_v1': ['Temperature'],
            
            # Thermal conductivity functions
            'CFUN_kCu_v1': ['Temperature'],  # May also use Magnetic Field and RRR
            'CFUN_kBrass_v1': ['Temperature'],
            'CFUN_kAl6061_v1': ['Temperature'],
            'CFUN_kAl1350_v1': ['Temperature'],
            'CFUN_kHast_v1': ['Temperature'],
            'CFUN_kAg_v1': ['Temperature'],
            'CFUN_kSteel_v2': ['Temperature'],
            
            # HTS superconductor critical current functions
            'CFUN_HTS_JcFit_Fujikura_v1': ['Temperature', 'Magnetic Field', 'Field Angle'],
            'CFUN_HTS_JcFit_THEVA_v1': ['Temperature', 'Magnetic Field', 'Field Angle'],
            'CFUN_HTS_JcFit_Benchmark_v1': ['Magnetic Field'],
        }
        
        # Check if the fit function name exists in our mapping
        if fit_name not in fit_parameters:
            available_fits = list(fit_parameters.keys())
            raise ValueError(
                f"Fit function '{fit_name}' is not recognized. "
                f"Available fit functions: {available_fits}"
            )
        
        return fit_parameters[fit_name]
    
    @staticmethod
    def validate_steam_material_call(fit_name: str, input_array: np.ndarray, expected_params: list[str] = None) -> bool:
        """
        Validate that a STEAM_materials call uses the correct input parameters.
        
        Args:
            fit_name (str): Name of the fit function
            input_array (np.ndarray): Input array being passed to STEAM_materials.evaluate()
            expected_params (list[str], optional): Expected parameters. If None, will be retrieved automatically.
            
        Returns:
            bool: True if validation passes
            
        Raises:
            ValueError: If validation fails with detailed error message
        """
        if expected_params is None:
            expected_params = MaterialProperties.get_fit_input_parameters(fit_name)
        
        expected_param_count = len(expected_params)
        actual_param_count = input_array.shape[0] if len(input_array.shape) > 1 else 1
        
        if actual_param_count != expected_param_count:
            raise ValueError(
                f"Parameter count mismatch for fit function '{fit_name}':\n"
                f"Expected {expected_param_count} parameters: {expected_params}\n"
                f"Got {actual_param_count} parameters in input array with shape {input_array.shape}\n"
                f"Make sure your input array has the correct structure: "
                f"np.vstack(({', '.join(expected_params)})) for multi-dimensional inputs"
            )
        
        return True
    
    @staticmethod
    def create_steam_material_input(fit_name: str, **kwargs) -> np.ndarray:
        """
        Create properly structured input array for STEAM_materials based on fit function requirements.
        
        Args:
            fit_name (str): Name of the fit function
            **kwargs: Parameter values (Temperature, Magnetic_Field, Field_Angle, etc.)
            
        Returns:
            np.ndarray: Properly structured input array for STEAM_materials.evaluate()
            
        Raises:
            ValueError: If required parameters are missing or unexpected parameters are provided
        """
        expected_params = MaterialProperties.get_fit_input_parameters(fit_name)
        
        # Map parameter names to more user-friendly kwargs
        param_mapping = {
            'Temperature': ['Temperature', 'T', 'temp'],
            'Magnetic Field': ['Magnetic_Field', 'Magnetic Field', 'B', 'magnetic_field'],
            'Field Angle': ['Field_Angle', 'Field Angle', 'angle', 'field_angle', 'deg']
        }
        
        input_arrays = []
        
        for param in expected_params:
            found = False
            param_value = None
            
            # Try to find the parameter under various names
            for possible_name in param_mapping.get(param, [param]):
                if possible_name in kwargs:
                    param_value = kwargs[possible_name]
                    found = True
                    break
            
            if not found:
                available_keys = list(kwargs.keys())
                raise ValueError(
                    f"Missing required parameter '{param}' for fit function '{fit_name}'.\n"
                    f"Required parameters: {expected_params}\n"
                    f"Available parameters: {available_keys}\n"
                    f"Hint: You can use these parameter names: {param_mapping.get(param, [param])}"
                )
            
            input_arrays.append(param_value)
        
        # Stack arrays if more than one parameter
        if len(input_arrays) == 1:
            input_array = np.array(input_arrays[0])
            # Ensure the array is at least 1D to avoid shape () issues
            if input_array.ndim == 0:
                input_array = input_array.reshape(1, 1)
            elif input_array.ndim == 1:
                input_array = input_array.reshape(1, -1)
            return input_array
        else:
            stacked = np.vstack(input_arrays)
            # Ensure proper 2D shape for multi-parameter case
            if stacked.ndim == 1:
                stacked = stacked.reshape(-1, 1)
            return stacked
    
    @staticmethod
    def density_selector(mat):
        """
        Select material density based on material name.
        
        Args:
            mat (str): Material name
            
        Returns:
            float: Material density in kg/m³
        """
        mat_rrr = re.findall(r'\d+', mat)  # Allows materials like "Cu20" -> mat = 'Cu' and rrr = '20.0'
        mat_rrr = np.float64(mat_rrr[0]) if len(mat_rrr) > 0 else 100.0
        mat = re.sub(r'\d+', "", mat).lower()
        if mat in ['cu', 'copper']:
            out = 8960.0
        elif mat == 'brass':
            out = 8730.0
        elif mat in ['al', 'aluminium', 'aluminum']:
            out = 2700.0
        elif mat in ['hastelloy', 'hast']:
            out = 8890.0
        elif mat in ['silver', 'ag']:
            out = 10490.0
        elif mat in ['ss', 'steel', 'stainless']:
            out = 7850.0
        else:
            supported_materials = [
                "'Cu', 'Copper'",
                "'Brass'",
                "'Al6061', 'Al5083', 'Al1350', or 'Al'",
                "'Hastelloy'",
                "'Silver'",
                "'Steel', 'SS', 'Stainless'"

            ]
            raise RuntimeError(
                f"Material '{mat}' is not in the density list.\n"
                f"Supported materials for density:\n  - " +
                "\n  - ".join(supported_materials)
            )
        return out

    def specific_heat_selector(self, mat, T, reshape=False):
        """
        Select material specific heat based on material name and temperature.

        Uses cached STEAM_materials objects for improved performance.

        Args:
            mat (str): Material name
            T (array): Temperature values
            reshape (bool): Whether to reshape temperature array

        Returns:
            array: Specific heat values
        """
        mat_rrr = re.findall(r'\d+', mat)  # Allows materials like "Cu20" -> mat = 'Cu' and rrr = '20.0'
        mat_rrr = np.float64(mat_rrr[0]) if len(mat_rrr) > 0 else 100.0
        mat = re.sub(r'\d+', "", mat).lower()
        if reshape:
            T = T.reshape((1, len(T)))
        if mat in ['cu', 'copper']:
            fit_name = 'CFUN_CvCu_v1'
            input_array = MaterialProperties.create_steam_material_input(fit_name, Temperature=T)
            max_points = input_array.shape[1]
            steam = self.get_cached_steam_material(fit_name, max_points=max_points)
            out = steam.evaluate(input_array)
        elif mat == 'al' and mat_rrr == 6061:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_CvAlAlloy_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat == 'al' and mat_rrr == 1350:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_CvAl_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat in ['hast', 'hastelloy']:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_CvHast_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat in ['steel', 'stainless steel', 'ss']:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_CvSteel_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat in ['ag', 'silver']:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_CvAg_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat in ['ss', 'steel', 'stainless']:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_CvSteel_v1', max_points=max_points)
            out = steam.evaluate(T)
        else:
            supported_materials = [
                "'cu' or 'copper'",
                "'al6061' or 'al1350'",
                "'hast' or 'hastelloy'",
                "'steel', 'stainless steel', or 'ss'",
                "'ag' or 'silver'"
            ]
            raise RuntimeError(
                f"Material '{mat}' is not in the specific heat list.\n"
                f"Supported materials for specific heat:\n  - " +
                "\n  - ".join(supported_materials)
            )
        return out

    def resistivity_selector(self, mat, T, B=None, reshape=False, set_B=None, B_min=None):
        """
        Select material resistivity based on material name, temperature, and magnetic field.

        Uses cached STEAM_materials objects for improved performance.

        Args:
            mat (str): Material name
            T (array): Temperature values
            B (array): Magnetic field values
            reshape (bool): Whether to reshape arrays
            set_B (float): Override magnetic field value
            B_min (float): Minimum magnetic field value

        Returns:
            array: Resistivity values
        """
        mat_rrr = re.findall(r'\d+', mat)  # Allows materials like "Cu20" -> mat = 'Cu' and rrr = '20.0'
        mat_rrr = np.float64(mat_rrr[0]) if len(mat_rrr) > 0 else 100.0
        mat = re.sub(r'\d+', "", mat).lower()
        if reshape:
            T = T.reshape((1, len(T)))
            B = B.reshape((1, len(B)))
        if mat in ['cu', 'copper']:
            if set_B is not None:
                B = set_B * np.ones_like(T)
            if B_min is not None:
                B = np.where(B<B_min, B_min, B)

            rrr = mat_rrr * np.ones_like(T)
            numpy2d = np.vstack((T, B, rrr))
            max_points = numpy2d.shape[1]
            steam = self.get_cached_steam_material('CFUN_rhoCu_NIST_v2', max_points=max_points)
            out = steam.evaluate(numpy2d)
            out = out if len(out) > 1 else out[0]
        elif mat == 'al' and mat_rrr == 6061:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_rhoAl6061_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat == 'al' and mat_rrr == 1350:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_rhoAl1350_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat in ['hastelloy', 'hast']:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_rhoHast_v2', max_points=max_points)
            out = steam.evaluate(T)
        elif mat == 'steel':
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_rhoSS_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat in ['ag', 'silver']:
            if set_B is not None:
                B = set_B * np.ones_like(T)
            rrr = 20. * np.ones(T.shape)
            rrr_ref = 295. * np.ones(T.shape)
            numpy2d = np.vstack((T, B, rrr, rrr_ref))
            max_points = numpy2d.shape[1]
            steam = self.get_cached_steam_material('CFUN_rhoAg_v1', max_points=max_points)
            out = steam.evaluate(numpy2d)
            out = out if len(out) > 1 else out[0]
        elif mat in ['ss', 'steel', 'stainless']:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_rhoSS_v1', max_points=max_points)
            out = steam.evaluate(T)
        else:
            supported_materials = [
                "'cu' or 'copper' (with optional RRR, e.g., 'Cu100')",
                "'al6061' or 'al1350'",
                "'hastelloy' or 'hast'",
                "'steel'",
                "'ag' or 'silver'"
            ]
            raise RuntimeError(
                f"Material '{mat+str(mat_rrr)}' is not in the resistivity list.\n"
                f"Supported materials for resistivity:\n  - " +
                "\n  - ".join(supported_materials)
            )
        return out

    def resistivity_selector_torch(self, mat, T, B=None, Normal_rho_FIT=None):
        """
        Torch-based resistivity selector for solver functions.
        
        Args:
            mat (str): Material name
            T (torch.Tensor): Temperature tensor
            B (torch.Tensor): Magnetic field tensor
            Normal_rho_FIT (dict): Precomputed interpolation functions
            
        Returns:
            torch.Tensor: Resistivity values
        """
        if B is None:
            B = torch.ones_like(T)*0.1
        if Normal_rho_FIT is None:
            raise ValueError("Normal_rho_FIT interpolation functions must be provided")
        out = Normal_rho_FIT[mat]((T, B))
        return out

    def conductivity_selector_torch(self, mat, T, B=None, Normal_k_FIT=None):
        """
        Torch-based conductivity selector for solver functions.
        
        Args:
            mat (str): Material name
            T (torch.Tensor): Temperature tensor
            B (torch.Tensor): Magnetic field tensor
            Normal_k_FIT (dict): Precomputed interpolation functions
            
        Returns:
            torch.Tensor: Thermal conductivity values
        """
        if B is None:
            B = torch.ones_like(T)*0.1
        if Normal_k_FIT is None:
            raise ValueError("Normal_k_FIT interpolation functions must be provided")
        out = Normal_k_FIT[mat]((T, B))
        return out

    def specific_heat_selector_torch(self, mat, T, B=None, Normal_Cv_FIT=None):
        """
        Torch-based specific heat selector for solver functions.
        
        Args:
            mat (str): Material name
            T (torch.Tensor): Temperature tensor
            B (torch.Tensor): Magnetic field tensor
            Normal_Cv_FIT (dict): Precomputed interpolation functions
            
        Returns:
            torch.Tensor: Specific heat values
        """
        if B is None:
            B = torch.ones_like(T)*0.1
        if Normal_Cv_FIT is None:
            raise ValueError("Normal_Cv_FIT interpolation functions must be provided")
        out = Normal_Cv_FIT[mat]((T, B))
        return out

    def conductivity_selector(self, mat, T, B=None, reshape=False, set_B=None):
        """
        Select material thermal conductivity based on material name and temperature.

        Uses cached STEAM_materials objects for improved performance.

        Args:
            mat (str): Material name
            T (array): Temperature values
            B (array): Magnetic field values
            reshape (bool): Whether to reshape arrays
            set_B (float): Override magnetic field value

        Returns:
            array: Thermal conductivity values
        """
        mat_rrr = re.findall(r'\d+', mat)  # Allows materials like "Cu20" -> mat = 'Cu' and rrr = '20.0'
        mat_rrr = np.float64(mat_rrr[0]) if len(mat_rrr) > 0 else 100.0
        mat = re.sub(r'\d+', "", mat).lower()
        if reshape:
            T = T.reshape((1, len(T)))
        if mat in ['cu', 'copper']:
            if set_B is not None:
                B = set_B * np.ones_like(T)
            numpy2d = np.vstack((T, 100 * np.ones_like(T), np.ones_like(T), np.ones_like(T))) #TODO check this function, RRR and B may be added.
            max_points = numpy2d.shape[1]
            steam = self.get_cached_steam_material('CFUN_kCu_v1', max_points=max_points)
            out = steam.evaluate(numpy2d)
            out = out if len(out) > 1 else out[0]
        elif mat == 'brass':
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_kBrass_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat == 'al' and mat_rrr == 6061:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_kAl6061_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat == 'al' and mat_rrr == 1350:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_kAl1350_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat in ['hastelloy', 'hast']:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_kHast_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat in ['ag', 'silver']:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_kAg_v1', max_points=max_points)
            out = steam.evaluate(T)
        elif mat in ['ss', 'steel', 'stainless']:
            max_points = T.shape[1]
            steam = self.get_cached_steam_material('CFUN_kSteel_v2', max_points=max_points)
            out = steam.evaluate(T)
        else:
            supported_materials = [
                "'cu' or 'copper' (with optional RRR, e.g., 'Cu100')",
                "'brass'",
                "'al6061' or 'al1350'",
                "'hastelloy' or 'hast'",
                "'ag' or 'silver'",
                "'steel' or 'ss'"
            ]
            raise RuntimeError(
                f"Material '{mat+str(mat_rrr)}' is not in the conductivity list.\n"
                f"Supported materials for thermal conductivity:\n  - " +
                "\n  - ".join(supported_materials)
            )
        return out
