"""
RUN FAIR ADVERSARIAL VALIDATION
================================
Usage:
python run_fair_adversarial_validation.py \
--model_path path/to/model.zip \
--num_episodes 100 \
--gpu \
--output_dir ./fair_test_results \
--seed # \
--start_test #
"""

import argparse
import sys
from pathlib import Path
import torch
import numpy as np

from gru_belief_policy_v3_agnostic import GRUAugmentedFeaturesExtractor
from gru_env_wrappers import GRUStateManager
from masked_dqn_policy import MaskedDQN, MaskedDQNPolicy
from stable_baselines3 import DQN

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent))

from fair_adversarial_validation_framework import FairAdversarialTestSuite
from mems_adaptive_inspection_env_curriculum_v5_SOFT_RESET_STABLE import (
    ResolutionAgnosticInspectionEnv,
    InspectionConfig,
    NaNSafetyWrapper
)

def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(
        description='Run fair adversarial validation on MEMS inspection model'
    )

    parser.add_argument(
        '--model_path',
        type=str,
        required=True,
        help='Path to trained model (.zip file)'
    )

    parser.add_argument(
        '--num_episodes',
        type=int,
        default=50,
        help='Number of episodes per test case (default: 50)'
    )

    parser.add_argument(
        '--output_dir',
        type=str,
        default='./fair_adversarial_results',
        help='Output directory for results'
    )

    parser.add_argument(
        '--grid_size',
        type=int,
        default=64,
        help='Environment grid size (default: 64)'
    )

    parser.add_argument(
        '--budget',
        type=int,
        default=3000,
        help='Inspection budget (default: 3000)'
    )

    parser.add_argument(
        '--gpu',
        action='store_true',
        help='Use GPU acceleration for environment'
    )

    parser.add_argument(
        '--seed',
        type=int,
        default=42,
        help='Random seed (default: 42)'
    )

    # Resume support
    parser.add_argument(
        '--start_test',
        type=int,
        default=1,
        help='Start from this test number (1-based). Use to resume after timeout/crash.'
    )

    return parser.parse_args()


def main():
    """Main execution"""
    args = parse_args()

    print("""
╔════════════════════════════════════════════════════════════════╗
║                                                                 ║
║  FAIR ADVERSARIAL VALIDATION FOR MEMS INSPECTION DQN           ║
║                                                                 ║
║  Testing model with realistic, interpretable scenarios         ║
║                                                                 ║
╚════════════════════════════════════════════════════════════════╝
""")

    # Set random seeds
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    # Load model
    print(f"\n{'='*80}")
    print(f"Loading Model from: {args.model_path}")
    print(f"{'='*80}")

    if not Path(args.model_path).exists():
        print(f"❌ Model file not found: {args.model_path}")
        sys.exit(1)

    custom_objects = {
        "features_extractor_class": GRUAugmentedFeaturesExtractor,
        "policy_class": MaskedDQNPolicy,
    }

    model = MaskedDQN.load(
        args.model_path, 
        custom_objects=custom_objects, 
        device="cuda" if args.gpu else "cpu"
    )
    print(f"✅ Model loaded successfully")

    # Create environment configuration
    print(f"\n{'='*80}")
    print(f"Creating Environment Configuration")
    print(f"{'='*80}")

    config = InspectionConfig(
        grid_size=args.grid_size,
        wafer_diameter_mm=300.0,
        die_size_mm=5.0,
        inspection_budget=args.budget,
        inspection_cost=1.0,
        defect_catch_value=100.0,
        false_alarm_penalty=2.0,
        miss_penalty=50.0,
        max_steps=None,
        prior_belief=0.1,
        belief_update_radius=3,
        belief_increase_rate=0.3,
        belief_decrease_rate=0.2,
        soft_reset=True,
        seed=args.seed
    )

    print(f"✅ Configuration created")
    print(f"   Grid size: {config.grid_size}x{config.grid_size}")
    print(f"   Budget: {config.inspection_budget}")
    print(f"   GPU acceleration: {args.gpu and torch.cuda.is_available()}")

    # Environment factory
    def env_factory():
        import copy
        fresh_config = copy.deepcopy(config)
        e = ResolutionAgnosticInspectionEnv(config=fresh_config, use_gpu=args.gpu)
        e = NaNSafetyWrapper(e)
        e = GRUStateManager(e, policy=model.policy)
        return e

    # Base env for cleanup
    env = env_factory()

    # Create test suite with resume support
    print(f"\n{'='*80}")
    print(f"Initializing Fair Test Suite")
    print(f"{'='*80}")

    test_suite = FairAdversarialTestSuite(
        model=model,
        env_factory=env_factory,
        output_dir=args.output_dir,
        start_test=args.start_test
    )

    print(f"Running {len(test_suite.test_cases)} fair tests")
    print(f"  - Production scenarios: {sum(1 for tc in test_suite.test_cases if tc.difficulty == 'production')}")
    print(f"  - Stress tests: {sum(1 for tc in test_suite.test_cases if tc.difficulty == 'stress')}")
    print(f"  - Extreme tests: {sum(1 for tc in test_suite.test_cases if tc.difficulty == 'extreme')}")
    if args.start_test > 1:
        print(f"  - Resuming from Test #{args.start_test}")

    # Run tests
    try:
        test_suite.run_all_tests(num_episodes_per_test=args.num_episodes)

        print(f"\n{'='*80}")
        print(f"✅ FAIR ADVERSARIAL VALIDATION COMPLETE")
        print(f"{'='*80}")
        print(f"Results saved to: {args.output_dir}")
        print(f"Review the files in {args.output_dir}")
        print(f"{'='*80}\n")

    except KeyboardInterrupt:
        print("\n\n⚠️ Validation interrupted by user")
        sys.exit(1)

    except Exception as e:
        print(f"\n\n❌ Error during validation: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

    finally:
        env.close()


if __name__ == "__main__":
    main()