import pandas as pd
import numpy as np
from typing import List, Dict, Optional
from utils.quantum_algorithms import QuantumInspiredOptimizer
from utils.data_loader import fetch_stock_data

class PortfolioManager:
    def __init__(self):
        self.quantum_optimizer = QuantumInspiredOptimizer()
        self.watchlists = {}  # User watchlists
        self.portfolios = {}  # User portfolios

    def create_watchlist(self, user_id: str, name: str) -> Dict:
        """Create a new watchlist for a user"""
        if user_id not in self.watchlists:
            self.watchlists[user_id] = {}
        
        self.watchlists[user_id][name] = []
        return {'status': 'success', 'message': f'Watchlist {name} created'}

    def add_to_watchlist(self, user_id: str, watchlist_name: str, symbol: str) -> Dict:
        """Add a symbol to a watchlist"""
        if user_id in self.watchlists and watchlist_name in self.watchlists[user_id]:
            if symbol not in self.watchlists[user_id][watchlist_name]:
                self.watchlists[user_id][watchlist_name].append(symbol)
            return {'status': 'success', 'message': f'Added {symbol} to watchlist'}
        return {'status': 'error', 'message': 'Watchlist not found'}

    def get_watchlist(self, user_id: str, watchlist_name: str) -> List[Dict]:
        """Get watchlist with current prices and analysis"""
        if user_id not in self.watchlists or watchlist_name not in self.watchlists[user_id]:
            return []

        watchlist_data = []
        for symbol in self.watchlists[user_id][watchlist_name]:
            data = fetch_stock_data(symbol, period='1d')
            if data is not None:
                current_price = data['Close'].iloc[-1]
                change = ((current_price - data['Open'].iloc[0]) / data['Open'].iloc[0]) * 100
                
                watchlist_data.append({
                    'symbol': symbol,
                    'current_price': current_price,
                    'change_percent': change,
                    'last_updated': data.index[-1].strftime('%Y-%m-%d %H:%M:%S')
                })

        return watchlist_data

    def optimize_portfolio(self, symbols: List[str], risk_tolerance: float = 0.5) -> Dict:
        """Optimize portfolio allocation using quantum-inspired algorithm"""
        # Fetch historical data for all symbols
        data = {}
        for symbol in symbols:
            hist_data = fetch_stock_data(symbol, period='1y')
            if hist_data is not None:
                data[symbol] = hist_data['Close']

        if not data:
            return {'status': 'error', 'message': 'No data available for optimization'}

        # Calculate returns
        returns = pd.DataFrame(data).pct_change().dropna()

        # Get optimal weights
        weights = self.quantum_optimizer.optimize_portfolio(returns, risk_tolerance)

        # Calculate portfolio metrics
        portfolio_return = sum(weights[symbol] * returns[symbol].mean() for symbol in symbols)
        portfolio_risk = np.sqrt(sum(sum(
            weights[s1] * weights[s2] * returns[s1].cov(returns[s2])
            for s2 in symbols) for s1 in symbols))

        return {
            'status': 'success',
            'allocation': weights,
            'metrics': {
                'expected_return': portfolio_return * 100,  # Convert to percentage
                'risk': portfolio_risk * 100,  # Convert to percentage
                'sharpe_ratio': portfolio_return / portfolio_risk if portfolio_risk > 0 else 0
            }
        }

    def analyze_portfolio(self, portfolio: Dict[str, float]) -> Dict:
        """Analyze current portfolio performance and suggest rebalancing"""
        symbols = list(portfolio.keys())
        current_weights = list(portfolio.values())
        
        # Get optimal weights
        optimal_allocation = self.optimize_portfolio(symbols)
        
        if optimal_allocation['status'] == 'error':
            return optimal_allocation

        # Compare current vs optimal allocation
        rebalancing_needed = any(
            abs(portfolio[symbol] - optimal_allocation['allocation'][symbol]) > 0.05
            for symbol in symbols
        )

        return {
            'status': 'success',
            'current_allocation': portfolio,
            'optimal_allocation': optimal_allocation['allocation'],
            'metrics': optimal_allocation['metrics'],
            'rebalancing_needed': rebalancing_needed,
            'rebalancing_suggestions': {
                symbol: optimal_allocation['allocation'][symbol] - portfolio[symbol]
                for symbol in symbols
            } if rebalancing_needed else {}
        }