import plotly.graph_objects as go
from plotly.subplots import make_subplots
import numpy as np

class ChartRenderer:
    def __init__(self):
        pass
    
    def render_price_chart(self, prices, actions=None, current_step=0):
        """Render price chart with actions"""
        fig = go.Figure()
        
        if not prices:
            # Return empty figure if no data
            fig.update_layout(
                title="Price Chart - No Data Available",
                xaxis_title="Time Step",
                yaxis_title="Price",
                height=300,
                template="plotly_white"
            )
            return fig
        
        # Add price line
        fig.add_trace(go.Scatter(
            x=list(range(len(prices))),
            y=prices,
            mode='lines',
            name='Price',
            line=dict(color='blue', width=2)
        ))
        
        # Add action markers if provided
        if actions and len(actions) == len(prices):
            buy_indices = [i for i, action in enumerate(actions) if action == 1]
            sell_indices = [i for i, action in enumerate(actions) if action == 2]
            close_indices = [i for i, action in enumerate(actions) if action == 3]
            
            if buy_indices:
                fig.add_trace(go.Scatter(
                    x=buy_indices,
                    y=[prices[i] for i in buy_indices],
                    mode='markers',
                    name='Buy',
                    marker=dict(color='green', size=10, symbol='triangle-up', 
                              line=dict(width=2, color='darkgreen'))
                ))
            
            if sell_indices:
                fig.add_trace(go.Scatter(
                    x=sell_indices,
                    y=[prices[i] for i in sell_indices],
                    mode='markers',
                    name='Sell',
                    marker=dict(color='red', size=10, symbol='triangle-down', 
                              line=dict(width=2, color='darkred'))
                ))
            
            if close_indices:
                fig.add_trace(go.Scatter(
                    x=close_indices,
                    y=[prices[i] for i in close_indices],
                    mode='markers',
                    name='Close',
                    marker=dict(color='orange', size=8, symbol='x', 
                              line=dict(width=2, color='darkorange'))
                ))
        
        fig.update_layout(
            title=f"Price Chart (Step: {current_step})",
            xaxis_title="Time Step",
            yaxis_title="Price",
            height=300,
            showlegend=True,
            template="plotly_white"
        )
        
        return fig
    
    def create_performance_chart(self, net_worth_history, reward_history, initial_balance):
        """Create portfolio performance chart"""
        fig = make_subplots(
            rows=2, cols=1,
            subplot_titles=['Portfolio Value Over Time', 'Step Rewards'],
            vertical_spacing=0.15
        )
        
        if not net_worth_history:
            fig.update_layout(title="No Data Available", height=400)
            return fig
        
        # Portfolio value
        fig.add_trace(go.Scatter(
            x=list(range(len(net_worth_history))),
            y=net_worth_history,
            mode='lines+markers',
            name='Net Worth',
            line=dict(color='green', width=3),
            marker=dict(size=4)
        ), row=1, col=1)
        
        # Add initial balance reference line
        fig.add_hline(y=initial_balance, line_dash="dash", 
                     line_color="red", annotation_text="Initial Balance",
                     row=1, col=1)
        
        # Rewards as bar chart
        if reward_history:
            fig.add_trace(go.Bar(
                x=list(range(len(reward_history))),
                y=reward_history,
                name='Reward',
                marker_color=['green' if r >= 0 else 'red' for r in reward_history],
                opacity=0.7
            ), row=2, col=1)
        
        fig.update_layout(height=500, showlegend=False, template="plotly_white")
        fig.update_yaxes(title_text="Value ($)", row=1, col=1)
        fig.update_yaxes(title_text="Reward", row=2, col=1)
        fig.update_xaxes(title_text="Step", row=2, col=1)
        
        return fig
    
    def create_action_distribution(self, actions):
        """Create action distribution pie chart"""
        fig = go.Figure()
        
        if not actions:
            fig.update_layout(title="No Actions Available", height=300)
            return fig
        
        action_names = ['Hold', 'Buy', 'Sell', 'Close']
        action_counts = [actions.count(i) for i in range(4)]
        
        colors = ['blue', 'green', 'red', 'orange']
        
        fig = go.Figure(data=[go.Pie(
            labels=action_names,
            values=action_counts,
            hole=.4,
            marker_colors=colors,
            textinfo='label+percent+value',
            hoverinfo='label+percent+value'
        )])
        
        fig.update_layout(
            title="Action Distribution",
            height=350,
            annotations=[dict(text='Actions', x=0.5, y=0.5, font_size=16, showarrow=False)],
            template="plotly_white"
        )
        
        return fig
    
    def create_training_progress(self, training_history):
        """Create training progress visualization"""
        if not training_history:
            fig = go.Figure()
            fig.update_layout(title="No Training Data Available", height=500)
            return fig
        
        episodes = [h['episode'] for h in training_history]
        rewards = [h['reward'] for h in training_history]
        net_worths = [h['net_worth'] for h in training_history]
        losses = [h.get('loss', 0) for h in training_history]
        
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=['Episode Rewards', 'Portfolio Value',
                          'Training Loss', 'Moving Average Reward (5)'],
            specs=[[{}, {}], [{}, {}]]
        )
        
        # Rewards
        fig.add_trace(go.Scatter(
            x=episodes, y=rewards, mode='lines+markers',
            name='Reward', line=dict(color='blue', width=2),
            marker=dict(size=4)
        ), row=1, col=1)
        
        # Portfolio value
        fig.add_trace(go.Scatter(
            x=episodes, y=net_worths, mode='lines+markers',
            name='Net Worth', line=dict(color='green', width=2),
            marker=dict(size=4)
        ), row=1, col=2)
        
        # Loss
        if any(loss > 0 for loss in losses):
            fig.add_trace(go.Scatter(
                x=episodes, y=losses, mode='lines+markers',
                name='Loss', line=dict(color='red', width=2),
                marker=dict(size=4)
            ), row=2, col=1)
        
        # Moving average reward
        if len(rewards) > 5:
            ma_rewards = []
            for i in range(len(rewards)):
                start_idx = max(0, i - 4)
                ma = np.mean(rewards[start_idx:i+1])
                ma_rewards.append(ma)
            
            fig.add_trace(go.Scatter(
                x=episodes, y=ma_rewards, mode='lines',
                name='MA Reward (5)', line=dict(color='orange', width=3, dash='dash')
            ), row=2, col=2)
        
        fig.update_layout(
            height=600, 
            showlegend=True, 
            title_text="Training Progress Over Episodes",
            template="plotly_white"
        )
        
        return fig