import boto3
import time
import os
import torch
from datetime import datetime


class SpotInstanceTrainer:
    """
    Production-ready training en Spot Instances con checkpointing automático.
    Handles interruptions gracefully y resume desde último checkpoint.
    """

    def __init__(self, checkpoint_dir="/mnt/efs/checkpoints"):
        self.ec2 = boto3.client('ec2', region_name='us-east-1')
        self.checkpoint_dir = checkpoint_dir
        os.makedirs(checkpoint_dir, exist_ok=True)

    def request_spot_instance(self, instance_type="p3.8xlarge", max_price="2.50"):
        """
        Request Spot Instance con fallback a On-Demand si no disponible.
        
        Args:
            instance_type: GPU instance type (p3.8xlarge = 4 V100 GPUs)
            max_price: Max bid price/hr (típico: 30-50% On-Demand price)
        
        Returns:
            instance_id: ID de instancia EC2 lanzada
        """
        # Launch template con user data para auto-setup
        user_data_script = """#!/bin/bash
# Install dependencies
pip install torch torchvision boto3

# Mount EFS para checkpoints compartidos
sudo mount -t nfs4 -o nfsvers=4.1 \\
  fs-12345.efs.us-east-1.amazonaws.com:/ /mnt/efs

# Download training script desde S3
aws s3 cp s3://my-ml-bucket/training_script.py /home/ubuntu/

# Start training con auto-resume
python /home/ubuntu/training_script.py --resume-from-checkpoint
"""

        response = self.ec2.request_spot_instances(
            SpotPrice=max_price,
            InstanceCount=1,
            Type='persistent',  # Re-launch si terminated
            LaunchSpecification={
                'ImageId': 'ami-0abcdef1234567890',  # Deep Learning AMI
                'InstanceType': instance_type,
                'KeyName': 'my-keypair',
                'UserData': user_data_script,
                'IamInstanceProfile': {
                    'Arn': 'arn:aws:iam::123456:instance-profile/ML-Role'
                },
                'BlockDeviceMappings': [{
                    'DeviceName': '/dev/sda1',
                    'Ebs': {
                        'VolumeSize': 100,
                        'VolumeType': 'gp3',
                        'DeleteOnTermination': True
                    }
                }],
                'SecurityGroupIds': ['sg-12345'],
                'SubnetId': 'subnet-12345'
            }
        )

        request_id = response['SpotInstanceRequests'][0]['SpotInstanceRequestId']
        print(f"Spot request created: {request_id}")

        # Wait for fulfillment
        while True:
            response = self.ec2.describe_spot_instance_requests(
                SpotInstanceRequestIds=[request_id]
            )
            status = response['SpotInstanceRequests'][0]['Status']['Code']
            
            if status == 'fulfilled':
                instance_id = response['SpotInstanceRequests'][0]['InstanceId']
                print(f"Spot instance launched: {instance_id}")
                return instance_id
            elif status in ['price-too-low', 'capacity-not-available']:
                print(f"Spot unavailable ({status}), falling back to On-Demand...")
                return self._launch_ondemand_fallback(instance_type)
            
            time.sleep(10)

    def save_checkpoint(self, model, optimizer, epoch, loss, checkpoint_name=None):
        """
        Save training checkpoint to EFS (shared across instances).
        
        Checkpoint incluye:
        - Model state_dict
        - Optimizer state_dict
        - Epoch number
        - Loss
        - Timestamp
        """
        if checkpoint_name is None:
            checkpoint_name = f"checkpoint_epoch_{epoch}.pt"
        
        checkpoint_path = os.path.join(self.checkpoint_dir, checkpoint_name)
        
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            'timestamp': datetime.now().isoformat()
        }

        # Atomic write (tmp file + rename para evitar corruption)
        tmp_path = checkpoint_path + '.tmp'
        torch.save(checkpoint, tmp_path)
        os.rename(tmp_path, checkpoint_path)
        
        print(f"Checkpoint saved: {checkpoint_path}")

        # Upload to S3 como backup adicional
        s3 = boto3.client('s3')
        s3.upload_file(
            checkpoint_path,
            'my-ml-bucket',
            f'checkpoints/{checkpoint_name}'
        )

    def load_latest_checkpoint(self, model, optimizer):
        """
        Load most recent checkpoint si existe (para resume training).
        
        Returns:
            epoch: Last completed epoch (or 0 if no checkpoint)
        """
        checkpoints = [
            f for f in os.listdir(self.checkpoint_dir) 
            if f.endswith('.pt')
        ]
        
        if not checkpoints:
            print("No checkpoints found, starting from scratch.")
            return 0

        # Sort by modification time, get latest
        latest = max(
            checkpoints,
            key=lambda f: os.path.getmtime(os.path.join(self.checkpoint_dir, f))
        )
        
        checkpoint_path = os.path.join(self.checkpoint_dir, latest)
        checkpoint = torch.load(checkpoint_path)
        
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        
        print(f"Resumed from checkpoint: {latest}, epoch {checkpoint['epoch']}")
        return checkpoint['epoch']


# USAGE EXAMPLE
if __name__ == "__main__":
    trainer = SpotInstanceTrainer()
    
    # Launch Spot instance
    instance_id = trainer.request_spot_instance(
        instance_type="p3.8xlarge",
        max_price="2.50"  # 90% discount vs $12/hr On-Demand
    )
    
    # Training loop con checkpointing (ejecuta en la instancia)
    model = MyModel()
    optimizer = torch.optim.Adam(model.parameters())
    
    # Resume from last checkpoint si existe
    start_epoch = trainer.load_latest_checkpoint(model, optimizer)
    
    for epoch in range(start_epoch, 100):
        # Training logic here
        loss = train_one_epoch(model, dataloader, optimizer)
        
        # Save checkpoint cada epoch (crítico para Spot)
        if epoch % 1 == 0:  # Cada epoch
            trainer.save_checkpoint(model, optimizer, epoch, loss)
        
        print(f"Epoch {epoch}, Loss: {loss}")

import pinecone
from sentence_transformers import SentenceTransformer
import numpy as np


# PINECONE SCALAR QUANTIZATION
pinecone.init(api_key="your-api-key", environment="us-west1-gcp")

# Create index CON quantization enabled
pinecone.create_index(
    name="rag-knowledge-base-quantized",
    dimension=768,  # all-MiniLM-L6-v2 embedding dimension
    metric="cosine",
    # SCALAR QUANTIZATION CONFIGURATION
    pod_type="p1.x1",
    metadata_config={
        "indexed": ["category", "source"]
    },
    # NEW: Quantization spec
    index_config={
        "quantization": {
            "type": "scalar",  # scalar, product, or binary
            "accuracy_target": 0.98  # 98% accuracy target (2% tolerance)
        }
    }
)


# WEAVIATE PRODUCT QUANTIZATION
import weaviate

client = weaviate.Client(
    url="https://your-cluster.weaviate.network",
    auth_client_secret=weaviate.AuthApiKey(api_key="your-api-key")
)

# Create class CON product quantization
class_obj = {
    "class": "Document",
    "vectorizer": "none",  # We provide embeddings
    "properties": [
        {"name": "content", "dataType": ["text"]},
        {"name": "category", "dataType": ["string"]},
    ],
    # PRODUCT QUANTIZATION CONFIG
    "vectorIndexConfig": {
        "distance": "cosine",
        "pq": {
            "enabled": True,
            "segments": 96,  # Number of segments (tuning parameter)
            "centroids": 256,  # Centroids per segment
            "trainingLimit": 100000,  # Training examples
            "encoder": {
                "type": "kmeans",
                "distribution": "log-normal"
            }
        }
    }
}

client.schema.create_class(class_obj)


# EXAMPLE: Insert vectors con quantization automática
model = SentenceTransformer('all-MiniLM-L6-v2')

documents = [
    "AWS SageMaker is a fully managed ML platform...",
    "Vector databases store high-dimensional embeddings...",
    # ... 10 million more documents
]

# Batch upsert (Pinecone example)
index = pinecone.Index("rag-knowledge-base-quantized")
batch_size = 100

for i in range(0, len(documents), batch_size):
    batch = documents[i:i+batch_size]
    embeddings = model.encode(batch)
    
    # Pinecone auto-applies quantization during upsert
    to_upsert = [
        (f"doc-{i+j}", embeddings[j].tolist(), {"text": doc})
        for j, doc in enumerate(batch)
    ]
    index.upsert(vectors=to_upsert)

print(f"Inserted {len(documents)} docs with quantization")


# QUERY con quantized index (transparente, mismo API)
query = "How does SageMaker work?"
query_embedding = model.encode([query])[0]

results = index.query(
    vector=query_embedding.tolist(),
    top_k=5,
    include_metadata=True
)

print(f"Top results: {results}")