import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import os
import logging
import torch.nn.functional as F
from typing import List, Tuple

# Configure logger
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class CNNModelPyTorch(nn.Module):
    """
    CNN model for trading signals
    Simplified version for RL training
    """
    def __init__(self, window_size: int, num_features: int, output_size: int, timeframes: List[str]):
        super(CNNModelPyTorch, self).__init__()
        
        self.window_size = window_size
        self.num_features = num_features
        self.output_size = output_size
        self.timeframes = timeframes
        
        # Device configuration
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        logger.info(f"Using device: {self.device}")
        
        # Build model
        self.build_model()
        
        # Initialize optimizer and scheduler
        self.optimizer = optim.Adam(self.parameters(), lr=0.001)
        self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer, mode='max', factor=0.5, patience=5, verbose=True
        )
        
        # Move model to device
        self.to(self.device)
        
    def build_model(self):
        """Build the CNN architecture"""
        # First Convolutional Layer
        self.conv1 = nn.Conv1d(
            in_channels=self.num_features * len(self.timeframes),
            out_channels=32,
            kernel_size=3,
            padding=1
        )
        self.bn1 = nn.BatchNorm1d(32)
        
        # Second Convolutional Layer
        self.conv2 = nn.Conv1d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm1d(64)
        
        # Third Convolutional Layer
        self.conv3 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm1d(128)
        
        # Calculate size after convolutions
        conv_out_size = self.window_size * 128
        
        # Fully connected layers
        self.fc1 = nn.Linear(conv_out_size, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, self.output_size)
        
        # Additional output for value estimation
        self.value_fc = nn.Linear(256, 1)
    
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Forward pass through the network"""
        # Ensure input is on the correct device
        x = x.to(self.device)
        
        # Reshape input: [batch, window_size, features] -> [batch, channels, window_size]
        batch_size = x.size(0)
        x = x.permute(0, 2, 1)
        
        # Convolutional layers
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        
        # Flatten
        x = x.view(batch_size, -1)
        
        # Fully connected layers
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        
        # Split into advantage and value streams
        advantage = self.fc3(x)
        value = self.value_fc(x)
        
        # Combine value and advantage
        q_values = value + (advantage - advantage.mean(dim=1, keepdim=True))
        
        return q_values, value
    
    def predict(self, X):
        """Make predictions"""
        self.eval()
        
        # Convert to tensor if not already
        if not isinstance(X, torch.Tensor):
            X_tensor = torch.tensor(X, dtype=torch.float32).to(self.device)
        else:
            X_tensor = X.to(self.device)
        
        with torch.no_grad():
            q_values, value = self(X_tensor)
            q_values_np = q_values.cpu().numpy()
            actions = np.argmax(q_values_np, axis=1)
            
        return actions, q_values_np
    
    def save(self, path: str):
        """Save model weights"""
        os.makedirs(os.path.dirname(path), exist_ok=True)
        torch.save(self.state_dict(), f"{path}.pt")
        logger.info(f"Model saved to {path}.pt")
    
    def load(self, path: str):
        """Load model weights"""
        self.load_state_dict(torch.load(f"{path}.pt", map_location=self.device))
        self.eval()
        logger.info(f"Model loaded from {path}.pt")