"""
Script to convert existing Keras model to a more compatible format.

Use this if you have a model saved with a newer TensorFlow/Keras version
that has compatibility issues when loading.
"""
import os
import sys
from tensorflow import keras
import warnings

warnings.filterwarnings('ignore', category=UserWarning)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


def convert_model(input_path: str, output_path: str = None):
    """
    Convert Keras model to a more compatible format.
    
    Args:
        input_path: Path to existing model file
        output_path: Path to save converted model (default: same as input with _compatible suffix)
    """
    if not os.path.exists(input_path):
        print(f"Error: Model file not found: {input_path}")
        return False
    
    if output_path is None:
        base, ext = os.path.splitext(input_path)
        output_path = f"{base}_compatible{ext}"
    
    print(f"Loading model from: {input_path}")
    try:
        # Try loading with different methods
        try:
            model = keras.models.load_model(input_path, compile=False)
            print("✓ Model loaded successfully")
        except Exception as e:
            print(f"✗ Error loading model: {e}")
            print("Trying alternative loading methods...")
            
            # Try with safe_mode=False (Keras 3.x)
            try:
                model = keras.models.load_model(input_path, compile=False, safe_mode=False)
                print("✓ Model loaded with safe_mode=False")
            except:
                # Try using tf.keras
                import tensorflow as tf
                model = tf.keras.models.load_model(input_path, compile=False)
                print("✓ Model loaded using tf.keras")
        
        print(f"\nSaving converted model to: {output_path}")
        
        # Save in compatible format
        try:
            model.save(
                output_path,
                save_format='keras',
                include_optimizer=False
            )
            print(f"✓ Model saved successfully (Keras format, no optimizer)")
        except TypeError:
            # If include_optimizer not supported
            model.save(output_path, save_format='keras')
            print(f"✓ Model saved successfully (Keras format)")
        
        print(f"\nConversion complete!")
        print(f"Original: {input_path}")
        print(f"Converted: {output_path}")
        print(f"\nYou can now replace the original model with the converted one:")
        print(f"  mv {output_path} {input_path}")
        
        return True
        
    except Exception as e:
        print(f"✗ Conversion failed: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: python convert_model.py <model_path> [output_path]")
        print("\nExample:")
        print("  python convert_model.py models/anomaly_autoencoder_cpu.keras")
        print("  python convert_model.py models/anomaly_autoencoder_cpu.keras models/anomaly_autoencoder_cpu_new.keras")
        sys.exit(1)
    
    input_path = sys.argv[1]
    output_path = sys.argv[2] if len(sys.argv) > 2 else None
    
    success = convert_model(input_path, output_path)
    sys.exit(0 if success else 1)