# data wrangling
import pandas as pd
import numpy as np 

# Import optional dependency checker
from aimodelshare.utils.optional_deps import check_optional

# --- ML FRAMEWORKS IMPORT SECTION ---
# Initialize to None to prevent NameErrors later if imports fail
sklearn = None
torch = None
xgboost = None
tf = None
keras = None
pyspark = None

# ml frameworks
try:
    import sklearn
    from sklearn.model_selection import GridSearchCV, RandomizedSearchCV
except:
    check_optional("sklearn", "Scikit-learn")

try:
    import torch
except:
    check_optional("torch", "PyTorch")

try:
    import xgboost
except:
    check_optional("xgboost", "XGBoost")

try:
    import tensorflow as tf
    import keras
except:
    check_optional("tensorflow", "TensorFlow/Keras")

try:
    import pyspark
    from pyspark.sql import SparkSession
    from pyspark.ml import PipelineModel, Model
    from pyspark.ml.tuning import CrossValidatorModel, TrainValidationSplitModel
    from onnxmltools import convert_sparkml
except:
    check_optional("pyspark", "PySpark")


# onnx modules
import onnx
import skl2onnx
from skl2onnx import convert_sklearn
# tf2onnx import is lazy-loaded to avoid requiring TensorFlow for non-TF workflows
_TF2ONNX_AVAILABLE = None
_tf2onnx_module = None
_tensorflow_module = None
try:
    from torch.onnx import export
except:
    pass
from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer
import importlib
import onnxmltools
import onnxruntime as rt
from skl2onnx.common.data_types import FloatTensorType

# aims modules
from aimodelshare.aws import run_function_on_lambda, get_aws_client
from aimodelshare.reproducibility import set_reproducibility_env
from pandas.io.formats.style import Styler

# os etc
import os
import ast
import tempfile
import json
import re
import pickle
import requests
import sys
import shutil
from pathlib import Path
from zipfile import ZipFile
import wget             
from copy import copy
import psutil
from pympler import asizeof
from IPython.display import display, HTML, SVG
import logging


import networkx as nx
import warnings
from pathlib import Path
import time
import signal

# scikeras imports keras which requires TensorFlow - lazy load it
try:
    from scikeras.wrappers import KerasClassifier, KerasRegressor
    _SCIKERAS_AVAILABLE = True
except ImportError:
    _SCIKERAS_AVAILABLE = False
    KerasClassifier = None
    KerasRegressor = None


logging.getLogger().setLevel(logging.ERROR)

def _check_tf2onnx_available():
    """Check if tf2onnx and TensorFlow are available, and load them if needed.
    
    Returns:
        tuple: (tf2onnx_module, tensorflow_module) on success
        
    Raises:
        RuntimeError: If TensorFlow or tf2onnx are not installed
    """
    global _TF2ONNX_AVAILABLE, _tf2onnx_module, _tensorflow_module
    
    if _TF2ONNX_AVAILABLE is None:
        try:
            import tf2onnx as tf2onnx_temp
            import tensorflow as tf_temp
            _tf2onnx_module = tf2onnx_temp
            _tensorflow_module = tf_temp
            _TF2ONNX_AVAILABLE = True
        except ImportError as e:
            _TF2ONNX_AVAILABLE = False
            raise RuntimeError(
                "TensorFlow and tf2onnx are required for Keras model conversion to ONNX. "
                "Please install them with: pip install tensorflow tf2onnx"
            ) from e
    
    if not _TF2ONNX_AVAILABLE:
        raise RuntimeError(
            "TensorFlow and tf2onnx are required for Keras model conversion to ONNX. "
            "Please install them with: pip install tensorflow tf2onnx"
        )
    
    return _tf2onnx_module, _tensorflow_module

def _extract_onnx_metadata(onnx_model, framework):
    '''Extracts model metadata from ONNX file.'''

    # get model graph
    graph = onnx_model.graph

    # initialize metadata dict
    metadata_onnx = {}

    def _get_shape(dims):
        return [d.dim_value if d.HasField("dim_value") else None for d in dims]
    
    input_dims = graph.input[0].type.tensor_type.shape.dim
    output_dims = graph.output[0].type.tensor_type.shape.dim
    
    metadata_onnx["input_shape"] = _get_shape(input_dims)
    metadata_onnx["output_shape"] = _get_shape(output_dims)
    
    # get layers and activations NEW
    # match layers and nodes and initalizers in sinle object
    # insert here....

    # get layers & activations 
    layer_nodes = ['MatMul', 'Gemm', 'Conv'] #add MaxPool, Transpose, Flatten
    activation_nodes = ['Relu', 'Softmax']

    layers = []
    activations = []

    for op_id, op in enumerate(graph.node):

        if op.op_type in layer_nodes:
            layers.append(op.op_type)
            #if op.op_type == 'MaxPool':
                #activations.append(None)
            
        if op.op_type in activation_nodes:
            activations.append(op.op_type)


    # get shapes and parameters
    layers_shapes = []
    layers_n_params = []

    if framework == 'keras':
        try:
            initializer = list(reversed(graph.initializer))
            for layer_id, layer in enumerate(initializer):
                if(len(layer.dims)>= 2):
                    layers_shapes.append(layer.dims[1])

                    try:
                        n_params = int(np.prod(layer.dims) + initializer[layer_id-1].dims)
                    except:
                        n_params = None

                    layers_n_params.append(n_params)
        except:
            pass
                

    elif framework == 'pytorch':
        try:
            initializer = graph.initializer
            for layer_id, layer in enumerate(initializer):
                if(len(layer.dims)>= 2):
                    layers_shapes.append(layer.dims[0])
                    n_params = int(np.prod(layer.dims) + initializer[layer_id-1].dims)
                    layers_n_params.append(n_params)
        except:
            pass


    # get model architecture stats
    model_architecture = {'layers_number': len(layers),
                      'layers_sequence': layers,
                      'layers_summary': {i:layers.count(i) for i in set(layers)},
                      'layers_n_params': layers_n_params,
                      'layers_shapes': layers_shapes,
                      'activations_sequence': activations,
                      'activations_summary': {i:activations.count(i) for i in set(activations)}
                      }

    metadata_onnx["model_architecture"] = model_architecture

    return metadata_onnx



def _misc_to_onnx(model, initial_types, transfer_learning=None,
                    deep_learning=None, task_type=None):

    # generate metadata dict 
    metadata = {}
    
    # placeholders, need to be generated elsewhere
    metadata['model_id'] = None
    metadata['data_id'] = None
    metadata['preprocessor_id'] = None
    
    try:
        # infer ml framework from function call
        if xgboost is not None and isinstance(model, (xgboost.XGBClassifier, xgboost.XGBRegressor)):
            metadata['ml_framework'] = 'xgboost'
            onx = onnxmltools.convert.convert_xgboost(model, initial_types=initial_types)
    except:
        pass

    # also integrate lightGBM 
    
    # get model type from model object
    model_type = str(model).split('(')[0]
    metadata['model_type'] = model_type
    
    # get transfer learning bool from user input
    metadata['transfer_learning'] = transfer_learning

    # get deep learning bool from user input
    metadata['deep_learning'] = deep_learning
    
    # get task type from user input
    metadata['task_type'] = task_type
    
    # placeholders, need to be inferred from data 
    metadata['target_distribution'] = None
    metadata['input_type'] = None
    metadata['input_shape'] = None
    metadata['input_dtypes'] = None        
    metadata['input_distribution'] = None
    
    # get model config dict from sklearn model object
    metadata['model_config'] = str(model.get_params())

    # get model state from sklearn model object
    metadata['model_state'] = None


    model_architecture = {}

    if hasattr(model, 'coef_'):
        model_architecture['layers_n_params'] = [len(model.coef_.flatten())]
    if hasattr(model, 'solver'):
        model_architecture['optimizer'] = model.solver

    metadata['model_architecture'] = str(model_architecture)

    metadata['memory_size'] = asizeof.asizeof(model)    


    # placeholder, needs evaluation engine
    metadata['eval_metrics'] = None  
    
    # add metadata from onnx object
    # metadata['metadata_onnx'] = str(_extract_onnx_metadata(onx, framework='sklearn'))
    metadata['metadata_onnx'] = None

    meta = onx.metadata_props.add()
    meta.key = 'model_metadata'
    meta.value = str(metadata)

    return onx



def _sklearn_to_onnx(model, initial_types=None, transfer_learning=None,
                    deep_learning=None, task_type=None):
    '''Extracts metadata from sklearn model object.'''
    
    # check whether this is a fitted sklearn model
    # sklearn.utils.validation.check_is_fitted(model)

    # deal with pipelines and parameter search 
    try:
        if sklearn is not None:
            if isinstance(model, (GridSearchCV, RandomizedSearchCV)):
                model = model.best_estimator_

            if isinstance(model, sklearn.pipeline.Pipeline):
                model = model.steps[-1][1]
    except:
        pass

    # fix ensemble voting models
    if all([hasattr(model, 'flatten_transform'),hasattr(model, 'voting')]):
      model.flatten_transform=False
    
    # convert to onnx
    if initial_types == None: 
        feature_count=model.n_features_in_
        initial_types = [('float_input', FloatTensorType([None, feature_count]))]

    onx = convert_sklearn(model, initial_types=initial_types,target_opset={'': 15, 'ai.onnx.ml': 2})
    
    ## set model ir_version to ensure sklearn opsets work properly
    onx.ir_version = 8
    
    # generate metadata dict 
    metadata = {}
    
    # placeholders, need to be generated elsewhere
    metadata['model_id'] = None
    metadata['data_id'] = None
    metadata['preprocessor_id'] = None
    
    # infer ml framework from function call
    metadata['ml_framework'] = 'sklearn'
    
    # get model type from model object
    model_type = str(model).split('(')[0]
    metadata['model_type'] = model_type
    
    # get transfer learning bool from user input
    metadata['transfer_learning'] = transfer_learning

    # get deep learning bool from user input
    metadata['deep_learning'] = deep_learning
    
    # get task type from user input
    metadata['task_type'] = task_type
    
    # placeholders, need to be inferred from data 
    metadata['target_distribution'] = None
    metadata['input_type'] = None
    metadata['input_shape'] = None
    metadata['input_dtypes'] = None        
    metadata['input_distribution'] = None
    
    # get model config dict from sklearn model object
    metadata['model_config'] = str(model.get_params())

    # get weights for pretrained models 
    temp_dir = tempfile.gettempdir()
    temp_path = os.path.join(temp_dir, 'temp_file_name')

    with open(temp_path, 'wb') as f:
        pickle.dump(model, f)

    with open(temp_path, "rb") as f:
        pkl_str = f.read()

    metadata['model_weights'] = pkl_str
    
    # get model state from sklearn model object
    metadata['model_state'] = None

    # get model architecture     
    if model_type == 'MLPClassifier' or model_type == 'MLPRegressor':

        if model_type == 'MLPClassifier':
            loss = 'log-loss'
        if model_type == 'MLPRegressor':
            loss = 'squared-loss'

        n_params = []
        layer_dims = [model.n_features_in_] + model.hidden_layer_sizes + [model.n_outputs_]
        for i in range(len(layer_dims)-1):
            n_params.append(layer_dims[i]*layer_dims[i+1] + layer_dims[i+1])

        # insert data into model architecture dict 
        model_architecture = {'layers_number': len(model.hidden_layer_sizes),
                              'layers_sequence': ['Dense']*len(model.hidden_layer_sizes),
                              'layers_summary': {'Dense': len(model.hidden_layer_sizes)},
                              'layers_n_params': n_params, #double check 
                              'layers_shapes': model.hidden_layer_sizes,
                              'activations_sequence': [model.activation]*len(model.hidden_layer_sizes),
                              'activations_summary': {model.activation: len(model.hidden_layer_sizes)},
                              'loss': loss,
                              'optimizer': model.solver
                             }

        metadata['model_architecture'] = str(model_architecture)


    else:
        model_architecture = {}

        if hasattr(model, 'coef_'):
            model_architecture['layers_n_params'] = [len(model.coef_.flatten())]
        if hasattr(model, 'solver'):
            model_architecture['optimizer'] = model.solver

        metadata['model_architecture'] = str(model_architecture)

    metadata['memory_size'] = asizeof.asizeof(model)    

    # placeholder, needs evaluation engine
    metadata['eval_metrics'] = None  
    
    # add metadata from onnx object
    # metadata['metadata_onnx'] = str(_extract_onnx_metadata(onx, framework='sklearn'))
    metadata['metadata_onnx'] = None

    meta = onx.metadata_props.add()
    meta.key = 'model_metadata'
    meta.value = str(metadata)

    return onx

def get_pyspark_model_files_paths(directory):
    # Get list of relative path of all model files

    # initializing empty file paths list
    file_paths = []
  
    for path in Path(directory).rglob('*'):
        if not path.is_dir():
            file_paths.append(path.relative_to(directory))

    # returning all file paths
    return file_paths

def _pyspark_to_onnx(model, initial_types, spark_session, 
                    transfer_learning=None, deep_learning=None, 
                    task_type=None):
    '''Extracts metadata from pyspark model object.'''

    try:
        if pyspark is None:
            raise("Error: Please install pyspark to enable pyspark features")
    except:
        raise("Error: Please install pyspark to enable pyspark features")

    # deal with pipelines and parameter search
    try:
        if isinstance(model, (TrainValidationSplitModel, CrossValidatorModel)):
            model = model.bestModel

        whole_model = copy(model)
        
        # Look for the last model in the pipeline
        if isinstance(model, PipelineModel):
            for t in model.stages:
                if isinstance(t, Model):
                    model = t
    except:
        pass

    # convert to onnx
    onx = convert_sparkml(whole_model, 'Pyspark model', initial_types, 
                          spark_session=spark_session)
            
    # generate metadata dict 
    metadata = {}
    
    # placeholders, need to be generated elsewhere
    metadata['model_id'] = None
    metadata['data_id'] = None
    metadata['preprocessor_id'] = None
    
    # infer ml framework from function call
    metadata['ml_framework'] = 'pyspark'
    
    # get model type from model object
    model_type = str(model).split(':')[0]
    metadata['model_type'] = model_type
    
    # get transfer learning bool from user input
    metadata['transfer_learning'] = transfer_learning

    # get deep learning bool from user input
    metadata['deep_learning'] = deep_learning
    
    # get task type from user input
    metadata['task_type'] = task_type
    
    # placeholders, need to be inferred from data 
    metadata['target_distribution'] = None
    metadata['input_type'] = None
    metadata['input_shape'] = None
    metadata['input_dtypes'] = None        
    metadata['input_distribution'] = None
    
    # get model config dict from pyspark model object
    model_config = {}
    try:
        for key, value in model.extractParamMap().items():
            model_config[key.name] = value
    except:
        pass
    metadata['model_config'] = str(model_config)

    # get weights for pretrained models 
    temp_dir = tempfile.gettempdir()
    temp_path = os.path.join(temp_dir, 'temp_pyspark_model')

    model.write().overwrite().save(temp_path)

    # calling function to get all file paths in the directory
    file_paths = get_pyspark_model_files_paths(temp_path)

    temp_path_zip = os.path.join(temp_dir, 'temp_pyspark_model.zip')
    with ZipFile(temp_path_zip, 'w') as zip:
        # writing each file one by one
        for file in file_paths:
            zip.write(os.path.join(temp_path, file), file)
    
    with open(temp_path_zip, "rb") as f:
        pkl_str = f.read()

    metadata['model_weights'] = pkl_str
    
    # clean up temp directory files for future runs
    try:
        shutil.rmtree(temp_path)
        os.remove(temp_path_zip)
    except:
        pass

    # get model state from sklearn model object
    metadata['model_state'] = None

    # get model architecture     
    if model_type == 'MultilayerPerceptronClassificationModel':

        #  https://spark.apache.org/docs/latest/ml-classification-regression.html#multilayer-perceptron-classifier
        loss = 'log-loss'
        hidden_layer_activation = 'sigmoid'
        output_layer_activation = 'softmax'

        n_params = []
        layer_dims = model.getLayers()
        hidden_layers = layer_dims[1:-1]
        for i in range(len(layer_dims)-1):
            n_params.append(layer_dims[i]*layer_dims[i+1] + layer_dims[i+1])

        # insert data into model architecture dict 
        model_architecture = {'layers_number': len(hidden_layers),
                              'layers_sequence': ['Dense']*len(hidden_layers),
                              'layers_summary': {'Dense': len(hidden_layers)},
                              'layers_n_params': n_params, #double check 
                              'layers_shapes': hidden_layers,
                              'activations_sequence': [hidden_layer_activation]*len(hidden_layers) + [output_layer_activation],
                              'activations_summary': {hidden_layer_activation: len(hidden_layers), output_layer_activation: 1},
                              'loss': loss,
                              'optimizer': model.getSolver()
                             }

        metadata['model_architecture'] = str(model_architecture)


    else:
        model_architecture = {}

        if hasattr(model, 'coefficients'):
            model_architecture['layers_n_params'] = [model.coefficients.size]
        if hasattr(model, 'getSolver') and callable(model.getSolver):
            model_architecture['optimizer'] = model.getSolver()

        metadata['model_architecture'] = str(model_architecture)

    metadata['memory_size'] = asizeof.asizeof(model)    

    # placeholder, needs evaluation engine
    metadata['eval_metrics'] = None  
    
    # add metadata from onnx object
    # metadata['metadata_onnx'] = str(_extract_onnx_metadata(onx, framework='sklearn'))
    metadata['metadata_onnx'] = None

    meta = onx.metadata_props.add()
    meta.key = 'model_metadata'
    meta.value = str(metadata)

    return onx

def _keras_to_onnx(model, transfer_learning=None,
                    deep_learning=None, task_type=None, epochs=None):
    '''Converts a Keras model to ONNX and extracts metadata.'''

    # Check and load tf2onnx and TensorFlow lazily (only when needed)
    tf2onnx, tf = _check_tf2onnx_available()
    
    import numpy as np
    import onnx
    import pickle
    import psutil
    import warnings
    from pympler import asizeof
    import logging
    import os
    import sys
    from contextlib import contextmanager

    # -- Helper to suppress tf2onnx stderr (NumPy warnings etc.)
    @contextmanager
    def suppress_stderr():
        with open(os.devnull, "w") as devnull:
            old_stderr = sys.stderr
            sys.stderr = devnull
            try:
                yield
            finally:
                sys.stderr = old_stderr

    # Reduce logging output
    tf2onnx_logger = logging.getLogger("tf2onnx")
    tf2onnx_logger.setLevel(logging.CRITICAL)

    # Unwrap scikeras, sklearn pipelines etc.
    try:
        from sklearn.model_selection import GridSearchCV, RandomizedSearchCV
        from sklearn.pipeline import Pipeline
        
        if isinstance(model, (GridSearchCV, RandomizedSearchCV)):
            model = model.best_estimator_
        if isinstance(model, Pipeline):
            model = model.steps[-1][1]
    except:
        pass

    try:
        from scikeras.wrappers import KerasClassifier, KerasRegressor
        if isinstance(model, (KerasClassifier, KerasRegressor)):
            model = model.model
    except:
        pass

    # Input signature
    input_shape = model.input_shape
    if isinstance(input_shape, list):
        input_shape = input_shape[0]
    input_signature = [tf.TensorSpec(input_shape, tf.float32, name="input")]

    # Wrap model in tf.function
    @tf.function(input_signature=input_signature)
    def model_fn(x):
        return model(x)

    concrete_func = model_fn

    # Convert to ONNX
    with suppress_stderr():
        onx_model, _ = tf2onnx.convert.from_function(
            concrete_func,
            input_signature=input_signature,
            opset=13,
            output_path=None
        )

    # Extract metadata
    metadata = {
        'model_id': None,
        'data_id': None,
        'preprocessor_id': None,
        'ml_framework': 'keras',
        'model_type': model.__class__.__name__,
        'transfer_learning': transfer_learning,
        'deep_learning': deep_learning,
        'task_type': task_type,
        'target_distribution': None,
        'input_type': None,
        'input_shape': input_shape,
        'input_dtypes': None,
        'input_distribution': None,
        'model_config': str(model.get_config()),
        'model_state': None,
        'eval_metrics': None,
        'model_graph': "",
        'metadata_onnx': None,
        'epochs': epochs
    }

    model_size = asizeof.asizeof(model.get_weights())
    mem = psutil.virtual_memory()

    if model_size > mem.available:
        warnings.warn(f"Model size ({model_size/1e6} MB) exceeds available memory.")
        metadata['model_weights'] = None
    else:
        metadata['model_weights'] = pickle.dumps(model.get_weights())

    # Extract architecture
    if not model.built: # add shape outputs if model not built
        try:
            model(tf.random.uniform([1] + list(input_shape[1:])))
        except Exception:
            pass  # fallback, don't crash conversion
            
    keras_layers = keras_unpack(model)
    
    
    from tensorflow.python.framework import tensor_shape  # <- place this at the top of your file
    
    layers = []
    layers_n_params = []
    layers_shapes = []
    activations = []
    
    for layer in keras_layers:
        # layer name
        layers.append(layer.__class__.__name__)
    
        # parameter count
        try:
            layers_n_params.append(layer.count_params())
        except:
            layers_n_params.append(0)
    
        # output shape (sanitized for JSON)
        shape = getattr(layer, 'output_shape', None)
    
        if isinstance(shape, tensor_shape.TensorShape):
            shape = shape.as_list()
        elif shape is not None:
            try:
                shape = list(shape)
            except:
                shape = str(shape)
        else:
            shape = None
    
        layers_shapes.append(shape)
    
        # activation
        if hasattr(layer, 'activation'):
            act = getattr(layer.activation, '__name__', None)
            if act:
                activations.append(act)

    optimizer = getattr(model.optimizer, '__class__', None)
    loss = getattr(model.loss, '__class__', None)

    model_architecture = {
        'layers_number': len(layers),
        'layers_sequence': layers,
        'layers_summary': {i: layers.count(i) for i in set(layers)},
        'layers_n_params': layers_n_params,
        'layers_shapes': layers_shapes,
        'activations_sequence': activations,
        'activations_summary': {i: activations.count(i) for i in set(activations)},
        'loss': loss.__name__ if loss else None,
        'optimizer': optimizer.__name__ if optimizer else None
    }

    metadata['model_architecture'] = str(model_architecture)
    metadata['model_summary'] = model_summary_keras(model).to_json()
    metadata['memory_size'] = model_size

    # Embed metadata in ONNX
    meta = onx_model.metadata_props.add()
    meta.key = 'model_metadata'
    meta.value = str(metadata)

    return onx_model



def _pytorch_to_onnx(model, model_input, transfer_learning=None, 
                    deep_learning=None, task_type=None, 
                    epochs=None):
    
    '''Extracts metadata from pytorch model object.'''

    # TODO check whether this is a fitted pytorch model
    # isinstance...

    # generate tempfile for onnx object 
    temp_dir = tempfile.gettempdir()
    temp_path = os.path.join(temp_dir, 'temp_file_name')

    # generate onnx file and save it to temporary path
    export(model, model_input, temp_path)
        #operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)

    # load onnx file from temporary path
    onx = onnx.load(temp_path)

    # generate metadata dict 
    metadata = {}

    # placeholders, need to be generated elsewhere
    metadata['model_id'] = None
    metadata['data_id'] = None
    metadata['preprocessor_id'] = None

    # infer ml framework from function call
    metadata['ml_framework'] = 'pytorch'

    # get model type from model object
    metadata['model_type'] = str(model.__class__).split('.')[-1].split("'")[0] + '()'

    # get transfer learning bool from user input
    metadata['transfer_learning'] = transfer_learning 

    # get deep learning bool from user input
    metadata['deep_learning'] = deep_learning

    # get task type from user input 
    metadata['task_type'] = task_type

    # placeholders, need to be inferred from data 
    metadata['target_distribution'] = None
    metadata['input_type'] = None
    metadata['input_shape'] = None
    metadata['input_dtypes'] = None        
    metadata['input_distribution'] = None

    # get model config dict from pytorch model object
    metadata['model_config'] = str(model.__dict__)

    # get model state from pytorch model object
    metadata['model_state'] = str(model.state_dict())


    name_list, layer_list, param_list, weight_list, activation_list = torch_metadata(model)

    model_summary_pd = pd.DataFrame({"Name": name_list,
    "Layer": layer_list,
    "Shape": weight_list,
    "Params": param_list,
    "Connect": None,
    "Activation": None})

    model_architecture = {'layers_number': len(layer_list),
                  'layers_sequence': layer_list,
                  'layers_summary': {i:layer_list.count(i) for i in set(layer_list)},
                  'layers_n_params': param_list,
                  'layers_shapes': weight_list,
                  'activations_sequence': activation_list,
                  'activations_summary': {i:activation_list.count(i) for i in set(activation_list)},
                  'loss': None,
                  'optimizer': None}

    metadata['model_architecture'] = str(model_architecture)

    metadata['model_summary'] = model_summary_pd.to_json()


    metadata['memory_size'] = asizeof.asizeof(model)    
    metadata['epochs'] = epochs

    # placeholder, needs evaluation engine
    metadata['eval_metrics'] = None

    # add metadata from onnx object
    # metadata['metadata_onnx'] = str(_extract_onnx_metadata(onx, framework='pytorch'))
    metadata['metadata_onnx'] = None

    # add metadata dict to onnx object
    meta = onx.metadata_props.add()
    meta.key = 'model_metadata'
    meta.value = str(metadata)
    

    return onx


def model_to_onnx(model, framework=None, model_input=None, initial_types=None,
                  transfer_learning=None, deep_learning=None, task_type=None, 
                  epochs=None, spark_session=None):
    
    '''Transforms sklearn, keras, or pytorch model object into ONNX format 
    and extracts model metadata dictionary. The model metadata dictionary 
    is saved in the ONNX file's metadata_props. 
    
    Parameters: 
    model: fitted sklearn, keras, or pytorch model object
    Specifies the model object that will be converted to ONNX. 
    
    framework: {"sklearn", "keras", "pytorch"}
    Specifies the machine learning framework of the model object.
    
    model_input: array_like, default=None
    Required when framework="pytorch".
    initial_types: initial types tuple, default=None
    Required when framework="sklearn".
      
    transfer_learning: bool, default=None
    Indicates whether transfer learning was used. 
    
    deep_learning: bool, default=None
    Indicates whether deep learning was used. 
    
    task_type: {"classification", "regression"}
    Indicates whether the model is a classification model or
    a regression model.
    
    Returns:
    ONNX object with model metadata saved in metadata props
    '''

    # if no framework was passed, extract framework 
    if model and framework==None:
        framework = model.__module__.split(".")[0]
        try:
            import torch
            if isinstance(model, torch.nn.Module):
                framework = "pytorch"
        except:
            pass

    # assert that framework exists
    frameworks = ['sklearn', 'keras', 'pytorch', 'xgboost', 'pyspark']
    assert framework in frameworks, \
    'Please choose "sklearn", "keras", "pytorch", "pyspark" or "xgboost".'
    
    # assert model input type THIS IS A PLACEHOLDER
    #if model_input is not None:
    #    assert isinstance(model_input, (list, pd.core.series.Series, np.ndarray, torch.Tensor)), \
    #    'Please format model input as XYZ.'
    
    # assert initialtypes 
    if initial_types != None:
        assert isinstance(initial_types[0][1], (skl2onnx.common.data_types.FloatTensorType)), \
        'Please use FloatTensorType as initial types.'
        
    # assert transfer_learning
    if transfer_learning != None:
        assert isinstance(transfer_learning, (bool)), \
        'Please pass boolean to indicate whether transfer learning was used.'
        
    # assert deep_learning
    if deep_learning != None:
        assert isinstance(deep_learning, (bool)), \
        'Please pass boolean to indicate whether deep learning was used.'
        
    # assert task_type
    if task_type != None:
        assert task_type in ['classification', 'regression'], \
        'Please specify task type as "classification" or "regression".'
    

    if framework == 'sklearn':
        onnx = _sklearn_to_onnx(model, initial_types=initial_types, 
                                transfer_learning=transfer_learning, 
                                deep_learning=deep_learning, 
                                task_type=task_type)
    elif framework == 'xgboost':
        onnx = _misc_to_onnx(model, initial_types=initial_types, 
                                transfer_learning=transfer_learning, 
                                deep_learning=deep_learning, 
                                task_type=task_type)
        
    elif framework == 'keras':
        onnx = _keras_to_onnx(model, transfer_learning=transfer_learning, 
                              deep_learning=deep_learning, 
                              task_type=task_type,
                              epochs=epochs)

        
    elif framework == 'pytorch':

        onnx = _pytorch_to_onnx(model, model_input=model_input,
                                transfer_learning=transfer_learning, 
                                deep_learning=deep_learning, 
                                task_type=task_type,
                                epochs=epochs)

    elif framework == 'pyspark':
        try:
            import pyspark
            from pyspark.sql import SparkSession
            from pyspark.ml import PipelineModel, Model
            from pyspark.ml.tuning import CrossValidatorModel, TrainValidationSplitModel
            from onnxmltools import convert_sparkml
        except:
            check_optional("pyspark", "PySpark")
        onnx = _pyspark_to_onnx(model, initial_types=initial_types, 
                                transfer_learning=transfer_learning, 
                                deep_learning=deep_learning, 
                                task_type=task_type,
                                spark_session=spark_session)

    try: 
        rt.InferenceSession(onnx.SerializeToString())   
    except Exception as e: 
        print(e)

    return onnx



def model_to_onnx_timed(model_filepath, force_onnx=False, timeout=60, model_input=None): 

    if not (model_filepath == None or isinstance(model_filepath, str) or isinstance(model_filepath, onnx.ModelProto)): 

        if force_onnx:
            try:
                import torch
                if isinstance(model_filepath, torch.nn.Module):
                    onnx_model = model_to_onnx(model_filepath, model_input=model_input)
            except:
                onnx_model = model_to_onnx(model_filepath)
            model_filepath = onnx_model

        else:


            try: # try needed because signal.alarm not available on windows
                # interrupt if onnx conversion is taking too long
                def timeout_handler(num, stack):
                    raise Exception("timeout")

                signal.alarm(timeout)
            except:
                pass


            try:
                try:
                    import torch
                    if isinstance(model_filepath, torch.nn.Module):
                        onnx_model = model_to_onnx(model_filepath, model_input=model_input)
                    else:
                        onnx_model = model_to_onnx(model_filepath)
                except:
                    onnx_model = model_to_onnx(model_filepath)
                model_filepath = onnx_model

            except:
                print("Timeout: Model to ONNX conversion is taking longer than expected. This can be the case for big models.")
                
                # Detect CI/testing environment for non-interactive fallback
                is_non_interactive = (
                    os.environ.get("PYTEST_CURRENT_TEST") is not None or
                    os.environ.get("AIMS_NON_INTERACTIVE") == "1"
                )
                
                if is_non_interactive:
                    # Auto-fallback to predictions-only in CI/testing environment
                    print("Non-interactive environment detected. Falling back to predictions-only submission.")
                    model_filepath = None
                else:
                    # Interactive prompt for manual runs
                    response = ''
                    while response not in {"1", "2"}:
                        response = input("Do you want to keep trying (1) or submit predictions only (2)? ")

                    if response == "1":
                        try:
                            import torch
                            if isinstance(model_filepath, torch.nn.Module):
                                onnx_model = model_to_onnx(model_filepath, model_input=model_input)
                            else:
                                onnx_model = model_to_onnx(model_filepath)
                        except Exception as e:
                            # Final fallback - if torch-specific handling failed, try generic conversion
                            # This handles cases where torch module detection fails but conversion might still work
                            warnings.warn(f"PyTorch-specific ONNX conversion failed ({e}), attempting generic conversion")
                            onnx_model = model_to_onnx(model_filepath, model_input=model_input)
                        model_filepath = onnx_model

                    elif response == "2":
                        model_filepath = None

            finally:
                print()
                try: # try needed because signal.alarm not available on windows
                    signal.alarm(0)
                except:
                    pass

    return model_filepath

def _get_metadata(onnx_model):
    '''Fetches previously extracted model metadata from ONNX object
    and returns model metadata dict.'''
    
    # double check this 
    #assert(isinstance(onnx_model, onnx.onnx_ml_pb2.ModelProto)), \
     #"Please pass a onnx model object."
    
    # Handle None input gracefully - always return a dict
    if onnx_model is None:
        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
            print("[DEBUG] _get_metadata: onnx_model is None, returning empty dict")
        return {}
    
    try: 
        onnx_meta = onnx_model.metadata_props

        onnx_meta_dict = {'model_metadata': ''}

        for i in onnx_meta:
            onnx_meta_dict[i.key] = i.value

        onnx_meta_dict = ast.literal_eval(onnx_meta_dict['model_metadata'])
        
        # Handle case where metadata is stored as a list instead of dict
        if isinstance(onnx_meta_dict, list):
            if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
                print(f"[DEBUG] _get_metadata: metadata is a list of length {len(onnx_meta_dict)}")
            if len(onnx_meta_dict) > 0 and isinstance(onnx_meta_dict[0], dict):
                onnx_meta_dict = onnx_meta_dict[0]
                if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
                    print("[DEBUG] _get_metadata: Extracted first dict from list")
            else:
                # Return empty dict if list doesn't contain valid dicts
                if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
                    print("[DEBUG] _get_metadata: List does not contain valid dicts, returning empty dict")
                return {}
        
        # Ensure we have a dict at this point
        if not isinstance(onnx_meta_dict, dict):
            if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
                print(f"[DEBUG] _get_metadata: Unexpected metadata type {type(onnx_meta_dict)}, returning empty dict")
            return {}
        
        #if onnx_meta_dict['model_config'] != None and \
        #onnx_meta_dict['ml_framework'] != 'pytorch':
        #    onnx_meta_dict['model_config'] = ast.literal_eval(onnx_meta_dict['model_config'])
        
        # Attempt to parse nested fields only if they are string representations of dicts
        if 'model_architecture' in onnx_meta_dict and onnx_meta_dict['model_architecture'] != None:
            try:
                if isinstance(onnx_meta_dict['model_architecture'], str):
                    onnx_meta_dict['model_architecture'] = ast.literal_eval(onnx_meta_dict['model_architecture'])
            except (ValueError, SyntaxError):
                # Keep as-is if parsing fails
                pass
            
        if 'model_config' in onnx_meta_dict and onnx_meta_dict['model_config'] != None:
            try:
                if isinstance(onnx_meta_dict['model_config'], str):
                    onnx_meta_dict['model_config'] = ast.literal_eval(onnx_meta_dict['model_config'])
            except (ValueError, SyntaxError):
                # Keep as-is if parsing fails
                pass
            
        if 'metadata_onnx' in onnx_meta_dict and onnx_meta_dict['metadata_onnx'] != None:
            try:
                if isinstance(onnx_meta_dict['metadata_onnx'], str):
                    onnx_meta_dict['metadata_onnx'] = ast.literal_eval(onnx_meta_dict['metadata_onnx'])
            except (ValueError, SyntaxError):
                # Keep as-is if parsing fails
                pass
        
        # onnx_meta_dict['model_image'] = onnx_to_image(onnx_model)

    except Exception as e:
    
        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
            print(f"[DEBUG] _get_metadata: Exception during metadata extraction: {e}")
        
        try:
            onnx_meta_dict = ast.literal_eval(onnx_meta_dict)
            # Handle list case in exception path as well
            if isinstance(onnx_meta_dict, list) and len(onnx_meta_dict) > 0 and isinstance(onnx_meta_dict[0], dict):
                onnx_meta_dict = onnx_meta_dict[0]
            elif not isinstance(onnx_meta_dict, dict):
                onnx_meta_dict = {}
        except:
            onnx_meta_dict = {}
    
    # Final safety check: ensure we always return a dict
    if not isinstance(onnx_meta_dict, dict):
        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
            print(f"[DEBUG] _get_metadata: Final check failed, returning empty dict instead of {type(onnx_meta_dict)}")
        return {}
        
    return onnx_meta_dict



def _get_leaderboard_data(onnx_model, eval_metrics=None):
    '''Extract leaderboard data from ONNX model or return defaults.
    
    This function performs single-pass normalization and safely handles:
    - None onnx_model (returns defaults)
    - Invalid metadata structures
    - Missing keys in metadata
    '''
    
    # Start with eval_metrics if provided, otherwise empty dict
    if eval_metrics is not None:
        metadata = dict(eval_metrics) if isinstance(eval_metrics, dict) else {}
    else:
        metadata = {}
    
    # Handle None onnx_model gracefully
    if onnx_model is None:
        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
            print("[DEBUG] _get_leaderboard_data: onnx_model is None, using default metadata")
        # Return metadata with safe defaults injected
        metadata['ml_framework'] = metadata.get('ml_framework', None)
        metadata['transfer_learning'] = metadata.get('transfer_learning', None)
        metadata['deep_learning'] = metadata.get('deep_learning', None)
        metadata['model_type'] = metadata.get('model_type', None)
        metadata['depth'] = metadata.get('depth', 0)
        metadata['num_params'] = metadata.get('num_params', 0)
        return metadata
        
    # Get metadata from ONNX - _get_metadata now always returns a dict
    metadata_raw = _get_metadata(onnx_model)
    
    if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
        print(f"[DEBUG] _get_leaderboard_data: metadata_raw type={type(metadata_raw)}, keys={list(metadata_raw.keys()) if isinstance(metadata_raw, dict) else 'N/A'}")
    
    # Single-pass normalization: ensure metadata_raw is a dict
    if not isinstance(metadata_raw, dict):
        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
            print(f"[DEBUG] _get_leaderboard_data: metadata_raw is not a dict (type={type(metadata_raw)}), using empty dict")
        metadata_raw = {}

    # get list of current layer types 
    layer_list_keras, activation_list_keras = _get_layer_names()
    layer_list_pytorch, activation_list_pytorch = _get_layer_names_pytorch()

    layer_list = list(set(layer_list_keras + layer_list_pytorch))
    activation_list =  list(set(activation_list_keras + activation_list_pytorch))

    # get general model info - use .get() for safety
    metadata['ml_framework'] = metadata_raw.get('ml_framework')
    metadata['transfer_learning'] = metadata_raw.get('transfer_learning')
    metadata['deep_learning'] = metadata_raw.get('deep_learning')
    metadata['model_type'] = metadata_raw.get('model_type')


    # get neural network metrics
    # Add isinstance check for model_architecture to prevent TypeError
    if (metadata_raw.get('ml_framework') in ['keras', 'pytorch'] or 
        metadata_raw.get('model_type') in ['MLPClassifier', 'MLPRegressor']) and \
        isinstance(metadata_raw.get('model_architecture'), dict):
        
        metadata['depth'] = metadata_raw['model_architecture'].get('layers_number', 0)
        metadata['num_params'] = sum(metadata_raw['model_architecture'].get('layers_n_params', []))

        for i in layer_list:
            layers_summary = metadata_raw['model_architecture'].get('layers_summary', {})
            if i in layers_summary:
                metadata[i.lower()+'_layers'] = layers_summary[i]
            elif i.lower()+'_layers' not in metadata.keys():
                metadata[i.lower()+'_layers'] = 0

        for i in activation_list:
            activations_summary = metadata_raw['model_architecture'].get('activations_summary', {})
            if i in activations_summary:
                if i.lower()+'_act' in metadata:
                    metadata[i.lower()+'_act'] += activations_summary[i]
                else:    
                    metadata[i.lower()+'_act'] = activations_summary[i]
            else:
                if i.lower()+'_act' not in metadata:
                    metadata[i.lower()+'_act'] = 0
                
        metadata['loss'] = metadata_raw['model_architecture'].get('loss')
        metadata['optimizer'] = metadata_raw['model_architecture'].get('optimizer')
        metadata['model_config'] = metadata_raw.get('model_config')
        metadata['epochs'] = metadata_raw.get('epochs')
        metadata['memory_size'] = metadata_raw.get('memory_size')

    # get sklearn & pyspark model metrics
    elif metadata_raw.get('ml_framework') in ['sklearn', 'xgboost', 'pyspark']:
        metadata['depth'] = 0

        try:
            if isinstance(metadata_raw.get('model_architecture'), dict):
                metadata['num_params'] = sum(metadata_raw['model_architecture'].get('layers_n_params', []))
            else:
                metadata['num_params'] = 0
        except:
            metadata['num_params'] = 0

        for i in layer_list:
            metadata[i.lower()+'_layers'] = 0

        for i in activation_list:
            metadata[i.lower()+'_act'] = 0

        metadata['loss'] = None

        try:
            if isinstance(metadata_raw.get('model_architecture'), dict):
                metadata['optimizer'] = metadata_raw['model_architecture'].get('optimizer')
            else:
                metadata['optimizer'] = None
        except:
            metadata['optimizer'] = None

        try:
            metadata['model_config'] = metadata_raw.get('model_config')
        except:
            metadata['model_config'] = None
    
    # Default handling for unknown frameworks
    else:
        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
            print(f"[DEBUG] _get_leaderboard_data: Unknown framework '{metadata_raw.get('ml_framework')}', using defaults")
        metadata.setdefault('depth', 0)
        metadata.setdefault('num_params', 0)
        for i in layer_list:
            metadata.setdefault(i.lower()+'_layers', 0)
        for i in activation_list:
            metadata.setdefault(i.lower()+'_act', 0)
    
    return metadata
    


def _model_summary(meta_dict, from_onnx=False):
    '''Creates model summary table from model metadata dict.'''
    import io
    
    assert(isinstance(meta_dict, dict)), \
    "Please pass valid metadata dict."
    
    assert('model_architecture' in meta_dict.keys()), \
    "Please make sure model architecture data is included."

    if from_onnx == True:
        model_summary = pd.read_json(io.StringIO(meta_dict['metadata_onnx']["model_summary"]))
    else:
        model_summary = pd.read_json(io.StringIO(meta_dict["model_summary"]))
        
    return model_summary



def onnx_to_image(model):
    '''Creates model graph image in pydot format.'''
    
    OP_STYLE = {
    'shape': 'box',
    'color': '#0F9D58',
    'style': 'filled',
    'fontcolor': '#FFFFFF'
    }
    
    pydot_graph = GetPydotGraph(
        model.graph,
        name=model.graph.name,
        rankdir='TB',
        node_producer=GetOpNodeProducer(
            embed_docstring=False,
            **OP_STYLE
        )
    )
    
    return pydot_graph


def inspect_model(apiurl, version=None, naming_convention = None, submission_type="competition"):
    if all(["username" in os.environ, 
           "password" in os.environ]):
        pass
    else:
        return print("'Inspect Model' unsuccessful. Please provide credentials with set_credentials().")

    post_dict = {"y_pred": [],
               "return_eval": "False",
               "return_y": "False",
               "inspect_model": "True",
               "version": version,
               "submission_type": submission_type
               }
    
    headers = { 'Content-Type':'application/json', 'authorizationToken': os.environ.get("AWS_TOKEN"),} 

    apiurl_eval=apiurl[:-1]+"eval"

    inspect_json = requests.post(apiurl_eval,headers=headers,data=json.dumps(post_dict)) 

    inspect_pd = pd.DataFrame(json.loads(inspect_json.text))

    return inspect_pd



def color_pal_assign(val, naming_convention=None):

    # find path of color mapping
    path =  Path(__file__).parent
    if naming_convention == "keras":
        col_map = pd.read_csv(path / "color_mappings/color_mapping_keras.csv")
    elif naming_convention == "pytorch":
        col_map = pd.read_csv(path / "color_mappings/color_mapping_pytorch.csv")

    # get color for layer key
    try:
        color = col_map[col_map.iloc[:,1] == val].iloc[:,2].values[0]
    except:
        color = "white"   

    return 'background: %s' % color


def stylize_model_comparison(comp_dict_out, naming_convention=None):

    for i in comp_dict_out.keys():

        if i == 'nn':

            df_styled = comp_dict_out['nn'].style.applymap(color_pal_assign, naming_convention=naming_convention)

            df_styled = df_styled.set_properties(**{'color': 'black'})

            df_styled = df_styled.set_caption('Model type: ' + 'Neural Network').set_table_styles([{'selector': 'caption',
                'props': [('color', 'white'), ('font-size', '18px')]}])

            df_styled = df_styled.set_properties(**{'color': 'black'})

            df_styled = df_styled.set_caption('Model type: ' + 'Neural Network').set_table_styles([{'selector': 'caption',
                'props': [('color', 'black'), ('font-size', '18px')]}])

            display(HTML(Styler.to_html(df_styled)))

        elif 'undefined' in i:

            version = i.split('_')[-1]

            df_styled = comp_dict_out[i].style

            df_styled = df_styled.set_caption("No metadata available for model "+ str(version)).set_table_styles([{'selector': 'caption',
                'props': [('color', 'black'), ('font-size', '18px')]}])

            display(HTML(Styler.to_html(df_styled)))
            print('\n\n')

        else:

            comp_pd = comp_dict_out[i]

            df_styled = comp_pd.style.apply(lambda x: ["background: tomato" if v != x.iloc[0] else "" for v in x], 
                axis = 1, subset=comp_pd.columns[1:])

            df_styled = df_styled.set_caption('Model type: ' + i).set_table_styles([{'selector': 'caption',
                'props': [('color', 'black'), ('font-size', '18px')]}])

            display(HTML(Styler.to_html(df_styled)))
            print('\n\n')



def compare_models(apiurl, version_list="None", 
    by_model_type=None, best_model=None, verbose=1, naming_convention=None, submission_type="competition"):
    if all(["username" in os.environ, 
           "password" in os.environ]):
        pass
    else:
        return print("'Inspect Model' unsuccessful. Please provide credentials with set_credentials().")

    post_dict = {"y_pred": [],
               "return_eval": "False",
               "return_y": "False",
               "inspect_model": "False",
               "version": "None", 
               "compare_models": "True",
               "version_list": version_list,
               "verbose": verbose, 
               "naming_convention": naming_convention,
               "submission_type": submission_type}
    
    headers = { 'Content-Type':'application/json', 'authorizationToken': os.environ.get("AWS_TOKEN"),} 

    apiurl_eval=apiurl[:-1]+"eval"

    compare_json = requests.post(apiurl_eval,headers=headers,data=json.dumps(post_dict)) 

    compare_dict = json.loads(compare_json.text)

    comp_dict_out = {i: pd.DataFrame(json.loads(compare_dict[i])) for i in compare_dict}

    return comp_dict_out


def _get_onnx_from_string(onnx_string):
    # generate tempfile for onnx object 
    temp_dir = tempfile.gettempdir()
    temp_path = os.path.join(temp_dir, 'temp_file_name')

    # save onnx to temporary path
    with open(temp_path, "wb") as f:
        f.write(onnx_string)

    # load onnx file from temporary path
    onx = onnx.load(temp_path)
    return onx

def _get_onnx_from_bucket(apiurl, aws_client, version=None):

    # generate name of onnx model in bucket
    onnx_model_name = "/onnx_model_v{version}.onnx".format(version = version)

    # Get bucket and model_id for user
    response, error = run_function_on_lambda(
        apiurl, **{"delete": "FALSE", "versionupdateget": "TRUE"}
    )
    if error is not None:
        raise error

    _, bucket, model_id = json.loads(response.content.decode("utf-8"))

    try:
        onnx_string = aws_client["client"].get_object(
            Bucket=bucket, Key=model_id + onnx_model_name
        )

        onnx_string = onnx_string['Body'].read()

    except Exception as err:
        raise err

    onx = _get_onnx_from_string(onnx_string)

    return onx



def instantiate_model(apiurl, version=None, trained=False, reproduce=False, submission_type="competition"):
    # Confirm that creds are loaded, print warning if not
    if all(["username" in os.environ, 
          "password" in os.environ]):
      pass
    else:
      return print("'Submit Model' unsuccessful. Please provide credentials with set_credentials().")

    post_dict = {
        "y_pred": [],
        "return_eval": "False",
        "return_y": "False",
        "inspect_model": "False",
        "version": "None", 
        "compare_models": "False",
        "version_list": "None",
        "get_leaderboard": "False",
        "instantiate_model": "True",
        "reproduce": reproduce,
        "trained": trained,
        "model_version": version,
        "submission_type": submission_type
    }

    headers = { 'Content-Type':'application/json', 'authorizationToken': os.environ.get("AWS_TOKEN"),} 

    apiurl_eval=apiurl[:-1]+"eval"

    resp = requests.post(apiurl_eval,headers=headers,data=json.dumps(post_dict)) 

    # Missing Check for response from Lambda. 
    try :
        resp.raise_for_status()
    except requests.exceptions.HTTPError :
        raise Exception(f"Error: Received {resp.status_code} from AWS, Please check if Model Version is correct.")


    resp_dict = json.loads(resp.text)

    if resp_dict['model_metadata'] == None:
        print("Model for this version doesn't exist or is not submitted by the author")
        return None

    if reproduce:
        if resp_dict['reproducibility_env'] != None:
            set_reproducibility_env(resp_dict['reproducibility_env'])
            print("Your reproducibility environment is successfully setup")
        else:
            print("Reproducibility environment is not found")

    print("Instantiate the model from metadata..")
    
    model_metadata = resp_dict['model_metadata']
    model_weight_url = resp_dict['model_weight_url']
    model_config = ast.literal_eval(model_metadata['model_config'])
    ml_framework = model_metadata['ml_framework']

    if ml_framework == 'sklearn':
        if trained == False or reproduce == True:
            model_type = model_metadata['model_type']
            model_class = model_from_string(model_type)
            model = model_class(**model_config)

        elif trained == True:
            model_pkl = None
            temp = tempfile.mkdtemp()
            temp_path = temp + "/" + "onnx_model_v{}.onnx".format(version)
            
            # Get leaderboard
            status = wget.download(model_weight_url, out=temp_path)
            onnx_model = onnx.load(temp_path)
            model_pkl = _get_metadata(onnx_model)['model_weights']
        
            temp_dir = tempfile.gettempdir()
            temp_path = os.path.join(temp_dir, 'temp_file_name')

            with open(temp_path, "wb") as f:
                f.write(model_pkl)

            with open(temp_path, 'rb') as f:
                model = pickle.load(f)

    if ml_framework == 'pyspark':
        try:
            if pyspark is None:
                raise("Error: Please install pyspark to enable pyspark features")
        except:
            raise("Error: Please install pyspark to enable pyspark features")

        if not trained or reproduce:
            print("Pyspark model can only be instantiated in trained mode.")
            print("Please rerun the function with proper parameters.")
            return None

        # pyspark model object is always trained. The unfitted / untrained one 
        # is the estimator and cannot be treated as model. 
        # Model is transformer and created by estimator
        model_pkl = None
        temp = tempfile.mkdtemp()
        temp_path = temp + "/" + "onnx_model_v{}.onnx".format(version)
        
        # Get leaderboard
        status = wget.download(model_weight_url, out=temp_path)
        onnx_model = onnx.load(temp_path)
        model_pkl = _get_metadata(onnx_model)['model_weights']

        temp_dir = tempfile.gettempdir()
        temp_path_zip = os.path.join(temp_dir, 'temp_pyspark_model.zip')
        temp_path = os.path.join(temp_dir, 'temp_pyspark_model')
        
        if not os.path.exists(temp_path):
            os.mkdir(temp_path)

        with open(temp_path_zip, "wb") as f:
            f.write(model_pkl)
        
        dirname_idx = -1
        with ZipFile(temp_path_zip, 'r') as zip_file:
            zip_file.extractall(temp_path)
        
        model_type = model_metadata['model_type']
        model_class = pyspark_model_from_string(model_type)
        # model_config is for the Estimator not the Transformer / Model
        model = model_class()

        # Need spark session and context to instantiate model object
        spark = SparkSession \
            .builder \
            .appName('Pyspark Model') \
            .getOrCreate()

        model = model.load(temp_path)

        # clean up temp directory files for future runs
        try:
            shutil.rmtree(temp_path)
            os.remove(temp_path_zip)
        except:
            pass

    if ml_framework == 'keras':
        if trained == False or reproduce == True:
            model = tf.keras.Sequential().from_config(model_config)

        elif trained == True:
            model_weights = None
            temp = tempfile.mkdtemp()
            temp_path = temp + "/" + "onnx_model_v{}.onnx".format(version)
        
            # Get leaderboard
            status = wget.download(model_weight_url, out=temp_path)
            onnx_model = onnx.load(temp_path)
            import pickle
            model_weights=pickle.loads(_get_metadata(onnx_model)['model_weights'])
            
            model = tf.keras.Sequential().from_config(model_config)

            model.set_weights(model_weights)

    print("Your model is successfully instantiated.")
    return model


def _get_layer_names():

    try:
        activation_list = [i for i in dir(tf.keras.activations)]
        activation_list = [i for i in activation_list if callable(getattr(tf.keras.activations, i))]
        activation_list = [i for i in activation_list if  not i.startswith("_")]
        activation_list.remove('deserialize')
        activation_list.remove('get')
        activation_list.remove('linear')
        activation_list = activation_list+['Activation', 'ReLU', 'Softmax', 'LeakyReLU', 'PReLU', 'ELU', 'ThresholdedReLU']


        layer_list = [i for i in dir(tf.keras.layers)]
        layer_list = [i for i in dir(tf.keras.layers) if callable(getattr(tf.keras.layers, i))]
        layer_list = [i for i in layer_list if not i.startswith("_")]
        layer_list = [i for i in layer_list if re.match('^[A-Z]', i)]
        layer_list = [i for i in layer_list if i.lower() not in [i.lower() for i in activation_list]]
    except:
        return [], []

    return layer_list, activation_list


def _get_layer_names_pytorch():

    activation_list = ['ELU', 'Hardshrink', 'Hardsigmoid', 'Hardtanh', 'Hardswish', 'LeakyReLU', 'LogSigmoid', 
                    'MultiheadAttention', 'PReLU', 'ReLU', 'ReLU6', 'RReLU', 'SELU', 'CELU', 'GELU', 'Sigmoid',
                    'SiLU', 'Mish', 'Softplus', 'Softshrink', 'Softsign', 'Tanh', 'Tanhshrink', 'Threshold',
                    'GLU', 'Softmin', 'Softmax', 'Softmax2d', 'LogSoftmax', 'AdaptiveLogSoftmaxWithLoss']
    try:
        import torch
        layer_list = [i for i in dir(torch.nn) if callable(getattr(torch.nn, i))]
        layer_list = [i for i in layer_list if not i in activation_list and not 'Loss' in i]
    except:
        layer_list=["no layers found"]
        pass


    return layer_list, activation_list


def _get_sklearn_modules():
    
    try:
        import sklearn

        sklearn_modules = ['ensemble', 'gaussian_process', 'isotonic',
                            'linear_model', 'mixture', 'multiclass', 'naive_bayes',
                            'neighbors', 'neural_network', 'svm', 'tree',
                            'discriminant_analysis', 'calibration']

        models_modules_dict = {}

        for i in sklearn_modules:
            models_list = [j for j in dir(eval('sklearn.'+i)) if callable(getattr(eval('sklearn.'+i), j))]
            models_list = [j for j in models_list if re.match('^[A-Z]', j)]

            for k in models_list: 
                models_modules_dict[k] = 'sklearn.'+i
    except:
        models_modules_dict = {}
    
    return models_modules_dict


def model_from_string(model_type):
    models_modules_dict = _get_sklearn_modules()
    try:
        module = models_modules_dict[model_type]
        model_class = getattr(importlib.import_module(module), model_type)
        return model_class
    except KeyError:
        # Return a placeholder class if estimator not found
        import warnings
        warnings.warn(f"Model type '{model_type}' not found in sklearn modules. Returning placeholder class.")
        
        # Create a minimal placeholder class that can be instantiated
        class PlaceholderModel:
            def __init__(self, **kwargs):
                self._model_type = model_type
                self._params = kwargs
            
            def get_params(self, deep=True):
                return self._params
            
            def __str__(self):
                return f"PlaceholderModel({self._model_type})"
            
            def __repr__(self):
                return f"PlaceholderModel({self._model_type})"
        
        return PlaceholderModel

def _get_pyspark_modules():
    try:
        if pyspark is None:
            raise Exception("Error: Please install pyspark to enable pyspark features")
    except:
        raise Exception("Error: Please install pyspark to enable pyspark features")

    pyspark_modules = ['ml', 'ml.feature', 'ml.classification', 'ml.clustering', 'ml.regression']

    models_modules_dict = {}

    for i in pyspark_modules:
        models_list = [j for j in dir(eval('pyspark.'+i)) if callable(getattr(eval('pyspark.'+i), j))]
        models_list = [j for j in models_list if re.match('^[A-Z]', j)]

        for k in models_list: 
            models_modules_dict[k] = 'pyspark.'+i
    
    return models_modules_dict


def pyspark_model_from_string(model_type):
    try:
        if pyspark is None:
            raise Exception("Error: Please install pyspark to enable pyspark features")
    except:
        raise Exception("Error: Please install pyspark to enable pyspark features")

    models_modules_dict = _get_pyspark_modules()
    module = models_modules_dict[model_type]
    model_class = getattr(importlib.import_module(module), model_type)
    return model_class


def print_y_stats(y_stats): 

  print("y_test example: ", y_stats['ytest_example'])
  print("y_test class labels", y_stats['class_labels'])
  print("y_test class balance", y_stats['class_balance'])
  print("y_test label dtypes", y_stats['label_dtypes'])


def inspect_y_test(apiurl, submission_type):

  # Confirm that creds are loaded, print warning if not
  if all(["username" in os.environ, 
          "password" in os.environ]):
      pass
  else:
      return print("'Submit Model' unsuccessful. Please provide credentials with set_credentials().")

  post_dict = {"y_pred": [],
               "return_eval": "False",
               "return_y": "True",
               "submission_type": submission_type}
    
  headers = { 'Content-Type':'application/json', 'authorizationToken': os.environ.get("AWS_TOKEN"),}

  apiurl_eval=apiurl[:-1]+"eval"

  y_stats = requests.post(apiurl_eval,headers=headers,data=json.dumps(post_dict)) 

  y_stats_dict = json.loads(y_stats.text)

  # print_y_stats(y_stats_dict)

  return y_stats_dict



def model_summary_keras(model):

    # extract model architecture metadata 
    layer_names = []
    layer_types = []
    layer_n_params = []
    layer_shapes = []
    layer_connect = []
    activations = []

    keras_layers = keras_unpack(model)

    for i in keras_layers: 

        try:
            layer_names.append(i.name)
        except:
            layer_names.append(None)

        try:
            layer_types.append(i.__class__.__name__)
        except:
            layer_types.append(None)

        try:
            layer_shapes.append(i.output_shape)
        except:
            layer_shapes.append(None)

        try:
            layer_n_params.append(i.count_params())
        except:
            layer_n_params.append(None)

        try:
            if isinstance(i.inbound_nodes[0].inbound_layers, list):
              layer_connect.append([x.name for x in i.inbound_nodes[0].inbound_layers])
            else: 
              layer_connect.append(i.inbound_nodes[0].inbound_layers.name)
        except:
            layer_connect.append(None)

        try: 
            activations.append(i.activation.__name__)
        except:
            activations.append(None)


    model_summary = pd.DataFrame({"Name": layer_names,
    "Layer": layer_types,
    "Shape": layer_shapes,
    "Params": layer_n_params,
    "Connect": layer_connect,
    "Activation": activations})

    return model_summary



def model_graph_keras(model):

    # extract model architecture metadata 
    layer_names = []
    layer_types = []
    layer_n_params = []
    layer_shapes = []
    layer_connect = []
    activations = []

    graph_nodes = []
    graph_edges = []


    for i in model.layers: 

        try:
            layer_name = i.name
        except:
            layer_name = None
        finally:
            layer_names.append(layer_name)


        try:
            layer_type = i.__class__.__name__
        except:
            layer_type = None
        finally:
            layer_types.append(layer_type)

        try:
            layer_shape = i.output_shape
        except:
            layer_shape = None
        finally:
            layer_shapes.append(layer_shape)


        try:
            layer_params = i.count_params()
        except:
            layer_params = None
        finally:
            layer_n_params.append(layer_params)


        try:
            if isinstance(i.inbound_nodes[0].inbound_layers, list):
              layer_input = [x.name for x in i.inbound_nodes[0].inbound_layers]
            else: 
              layer_input = i.inbound_nodes[0].inbound_layers.name
        except:
            layer_connect = None
        finally:
            layer_connect.append(layer_input)


        try: 
            activation = i.activation.__name__
        except:
            activation = None
        finally:
            activations.append(activation)

        layer_color = color_pal_assign(layer_type)
        layer_color =  layer_color.split(' ')[-1]


        graph_nodes.append((layer_name, {"label": layer_type + '\n' + str(layer_shape),
                                                 "URL": "https://keras.io/search.html?query="+layer_type.lower(),
                                                 "color": layer_color,
                                                 "style": "bold",
                                            "Name": layer_name,
                                            "Layer": layer_type,
                                            "Shape": layer_shape,
                                            "Params": layer_params,
                                            "Activation": activation}))
        
        if isinstance(layer_input, list):
            for i in layer_input:
                graph_edges.append((i, layer_name))
        else:
            graph_edges.append((layer_input, layer_name))

    G = nx.DiGraph()
    G.add_nodes_from(graph_nodes)
    G.add_edges_from(graph_edges)

    G_pydot = nx.drawing.nx_pydot.to_pydot(G)

    return G_pydot


def plot_keras(model):

    G =  model_graph_keras(model)

    display(SVG(G.create_svg()))



def torch_unpack(model):
    
    layers = []
    keys = []
    
    for key, module in model._modules.items():
                
        if len(module._modules):
            
            layers_out, keys_out = torch_unpack(module)
            
            layers = layers + layers_out
            keys = keys + keys_out
            
            
        else:
            
            layers.append(module)
            keys.append(key)
            
    return layers, keys

    
def keras_unpack(model):
    layers = []
    try:
        for module in model.layers:
            if isinstance(module, (tf.keras.Model, tf.keras.Sequential)):
                layers += keras_unpack(module)
            else:
                layers.append(module)
    except:
        pass
    return layers


def torch_metadata(model):

    name_list_out = []
    layer_list = []
    param_list = []
    weight_list = []
    activation_list = []
    
    try:
        layers, name_list = torch_unpack(model)

        layer_names, activation_names = _get_layer_names_pytorch()

        for module, name in zip(layers, name_list):

            module_name = module._get_name()


            if module_name in layer_names:

                    name_list_out.append(name)

                    layer_list.append(module_name)

                    params = sum([np.prod(p.size()) for p in module.parameters()])
                    param_list.append(params)

                    weights = tuple([tuple(p.size()) for p in module.parameters()])
                    weight_list.append(weights)

            if module_name in activation_names: 

                    activation_list.append(module_name)
    except:
        pass

    return name_list_out, layer_list, param_list, weight_list, activation_list


def layer_mapping(direction='torch_to_keras', activation=False):

    torch_keras = {'AdaptiveAvgPool1d': 'AvgPool1D',
    'AdaptiveAvgPool2d': 'AvgPool2D',
    'AdaptiveAvgPool3d': 'AvgPool3D',
    'AdaptiveMaxPool1d': 'MaxPool1D',
    'AdaptiveMaxPool2d': 'MaxPool2D',
    'AdaptiveMaxPool3d': 'MaxPool3D',
    'AlphaDropout': None,
    'AvgPool1d': 'AvgPool1D',
    'AvgPool2d': 'AvgPool2D',
    'AvgPool3d': 'AvgPool3D',
    'BatchNorm1d': 'BatchNormalization',
    'BatchNorm2d': 'BatchNormalization',
    'BatchNorm3d': 'BatchNormalization',
    'Bilinear': None,
    'ConstantPad1d': None,
    'ConstantPad2d': None,
    'ConstantPad3d': None,
    'Container': None,
    'Conv1d': 'Conv1D',
    'Conv2d': 'Conv2D',
    'Conv3d': 'Conv3D',
    'ConvTranspose1d': 'Conv1DTranspose',
    'ConvTranspose2d': 'Conv2DTranspose',
    'ConvTranspose3d': 'Conv3DTranspose',
    'CosineSimilarity': None,
    'CrossMapLRN2d': None,
    'DataParallel': None,
    'Dropout': 'Dropout',
    'Dropout2d': 'Dropout',
    'Dropout3d': 'Dropout',
    'Embedding': 'Embedding',
    'EmbeddingBag': 'Embedding',
    'FeatureAlphaDropout': None,
    'Flatten': 'Flatten',
    'Fold': None,
    'FractionalMaxPool2d': "MaxPool2D",
    'FractionalMaxPool3d': "MaxPool3D",
    'GRU': 'GRU',
    'GRUCell': 'GRUCell',
    'GroupNorm': None,
    'Identity': None,
    'InstanceNorm1d': None,
    'InstanceNorm2d': None,
    'InstanceNorm3d': None,
    'LPPool1d': None,
    'LPPool2d': None,
    'LSTM': 'LSTM',
    'LSTMCell': 'LSTMCell',
    'LayerNorm': None,
    'Linear': 'Dense',
    'LocalResponseNorm': None,
    'MaxPool1d': 'MaxPool1D',
    'MaxPool2d': 'MaxPool2D',
    'MaxPool3d': 'MaxPool3D',
    'MaxUnpool1d': None,
    'MaxUnpool2d': None,
    'MaxUnpool3d': None,
    'Module': None,
    'ModuleDict': None,
    'ModuleList': None,
    'PairwiseDistance': None,
    'Parameter': None,
    'ParameterDict': None,
    'ParameterList': None,
    'PixelShuffle': None,
    'RNN': 'RNN',
    'RNNBase': None,
    'RNNCell': None,
    'RNNCellBase': None,
    'ReflectionPad1d': None,
    'ReflectionPad2d': None,
    'ReplicationPad1d': None,
    'ReplicationPad2d': None,
    'ReplicationPad3d': None,
    'Sequential': None,
    'SyncBatchNorm': None,
    'Transformer': None,
    'TransformerDecoder': None,
    'TransformerDecoderLayer': None,
    'TransformerEncoder': None,
    'TransformerEncoderLayer': None,
    'Unfold': None,
    'Upsample': 'UpSampling1D',
    'UpsamplingBilinear2d': 'UpSampling2D',
    'UpsamplingNearest2d': 'UpSampling2D',
    'ZeroPad2d': 'ZeroPadding2D'}

    keras_torch = {'AbstractRNNCell': None,
    'Activation': None,
    'ActivityRegularization': None,
    'Add': None,
    'AdditiveAttention': None,
    'AlphaDropout': None,
    'Attention': None,
    'Average': None,
    'AveragePooling1D': 'AvgPool1d',
    'AveragePooling2D': 'AvgPool2d',
    'AveragePooling3D': 'AvgPool3d',
    'AvgPool1D': 'AvgPool1d',
    'AvgPool2D': 'AvgPool2d',
    'AvgPool3D': 'AvgPool3d',
    'BatchNormalization': None,
    'Bidirectional': None,
    'Concatenate': None,
    'Conv1D': 'Conv1d',
    'Conv1DTranspose': 'ConvTranspose1d',
    'Conv2D': 'Conv2d',
    'Conv2DTranspose':  'ConvTranspose2d',
    'Conv3D': 'Conv3d',
    'Conv3DTranspose':  'ConvTranspose3d',
    'ConvLSTM2D': None,
    'Convolution1D': None,
    'Convolution1DTranspose': None,
    'Convolution2D': None,
    'Convolution2DTranspose': None,
    'Convolution3D': None,
    'Convolution3DTranspose': None,
    'Cropping1D': None,
    'Cropping2D': None,
    'Cropping3D': None,
    'Dense': 'Linear',
    'DenseFeatures': None,
    'DepthwiseConv2D': None,
    'Dot': None,
    'Dropout': 'Dropout',
    'Embedding': 'Embedding',
    'Flatten': 'Flatten',
    'GRU': 'GRU',
    'GRUCell': 'GRUCell',
    'GaussianDropout': None,
    'GaussianNoise': None,
    'GlobalAveragePooling1D': None,
    'GlobalAveragePooling2D': None,
    'GlobalAveragePooling3D': None,
    'GlobalAvgPool1D': None,
    'GlobalAvgPool2D': None,
    'GlobalAvgPool3D': None,
    'GlobalMaxPool1D': None,
    'GlobalMaxPool2D': None,
    'GlobalMaxPool3D': None,
    'GlobalMaxPooling1D': None,
    'GlobalMaxPooling2D': None,
    'GlobalMaxPooling3D': None,
    'Input': None,
    'InputLayer': None,
    'InputSpec': None,
    'LSTM': 'LSTM',
    'LSTMCell': 'LSTMCell',
    'Lambda': None,
    'Layer': None,
    'LayerNormalization': None,
    'LocallyConnected1D': None,
    'LocallyConnected2D': None,
    'Masking': None,
    'MaxPool1D': 'MaxPool1d',
    'MaxPool2D': 'MaxPool2d',
    'MaxPool3D': 'MaxPool3d',
    'MaxPooling1D': 'MaxPool1d',
    'MaxPooling2D': 'MaxPool2d',
    'MaxPooling3D': 'MaxPool3d',
    'Maximum': None,
    'Minimum': None,
    'MultiHeadAttention': None,
    'Multiply': None,
    'Permute': None,
    'RNN': 'RNN',
    'RepeatVector': None,
    'Reshape': None,
    'SeparableConv1D': None,
    'SeparableConv2D': None,
    'SeparableConvolution1D': None,
    'SeparableConvolution2D': None,
    'SimpleRNN': None,
    'SimpleRNNCell': None,
    'SpatialDropout1D': None,
    'SpatialDropout2D': None,
    'SpatialDropout3D': None,
    'StackedRNNCells': None,
    'Subtract': None,
    'TimeDistributed': None,
    'UpSampling1D': 'Upsample',
    'UpSampling2D': None,
    'UpSampling3D': None,
    'Wrapper': None,
    'ZeroPadding1D': None,
    'ZeroPadding2D': 'ZeroPad2d',
    'ZeroPadding3D': None}

    torch_keras_act = {
    'AdaptiveLogSoftmaxWithLoss': None,
    'CELU': None,
    'ELU': 'elu',
    'GELU': 'gelu',
    'GLU': None,
    'Hardshrink': None,
    'Hardsigmoid': 'hard_sigmoid',
    'Hardswish': None,
    'Hardtanh': None,
    'LeakyReLU': 'LeakyReLU',
    'LogSigmoid': None,
    'LogSoftmax': None,
    'Mish': None,
    'MultiheadAttention': None,
    'PReLU': 'PReLU',
    'RReLU': None,
    'ReLU': 'relu',
    'ReLU6': 'relu',
    'SELU': 'selu',
    'SiLU': 'swish',
    'Sigmoid': 'sigmoid',
    'Softmax': 'softmax',
    'Softmax2d': None,
    'Softmin': None,
    'Softplus': 'softplus',
    'Softshrink': None,
    'Softsign': 'softsign',
    'Tanh': 'tanh',
    'Tanhshrink': None,
    'Threshold': None}

    keras_torch_act = {
    'ELU': 'ELU',
    'LeakyReLU': 'LeakyReLU',
    'PReLU': 'PReLU',
    'ReLU': 'ReLU',
    'Softmax': 'Softmax',
    'ThresholdedReLU': None,
    'elu': 'ELU',
    'exponential': None,
    'gelu': 'GELU',
    'hard_sigmoid': 'Hardsigmoid',
    'relu': 'ReLU',
    'selu': 'SELU',
    'serialize': None,
    'sigmoid': 'Sigmoid',
    'softmax': 'Softmax',
    'softplus': 'Softplus',
    'softsign': 'Softsign',
    'swish': 'SiLU',
    'tanh': 'Tanh'}


    if direction == 'torch_to_keras' and activation:

        return torch_keras_act

    elif direction == 'kreas_to_torch' and not activation:

        return keras_torch_act

    elif direction == 'torch_to_keras':

        return torch_keras

    elif direction == 'keras_to_torch': 

        return keras_torch


def rename_layers(in_layers, direction="torch_to_keras", activation=False):

  mapping_dict = layer_mapping(direction=direction, activation=activation)

  out_layers = []

  for i in in_layers:

    layer_name_temp = mapping_dict.get(i, None)

    if layer_name_temp == None:
      out_layers.append(i)
    else:
      out_layers.append(layer_name_temp)

  return out_layers

