"""
시각화 모듈
평가 결과를 시각화하고 대시보드 및 리포트를 생성합니다.
"""

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.offline as pyo
from typing import Dict, List, Optional, Tuple
import os
import json
from datetime import datetime
import trimesh
import open3d as o3d


class Visualizer:
    """시각화를 담당하는 클래스"""
    
    def __init__(self, output_dir: str = "results"):
        """
        시각화 도구 초기화
        
        Args:
            output_dir (str): 출력 디렉토리
        """
        self.output_dir = output_dir
        self.images_dir = os.path.join(output_dir, "images")
        self.reports_dir = os.path.join(output_dir, "reports")
        
        # 디렉토리 생성
        os.makedirs(self.images_dir, exist_ok=True)
        os.makedirs(self.reports_dir, exist_ok=True)
        
        # matplotlib 스타일 설정
        plt.style.use('seaborn-v0_8')
        sns.set_palette("husl")
    
    def plot_metrics_comparison(self, results: Dict, save_path: Optional[str] = None) -> None:
        """
        평가 지표 비교 차트를 생성합니다.
        
        Args:
            results (Dict): 평가 결과
            save_path (Optional[str]): 저장 경로
        """
        metrics = results['metrics']
        score_details = results['score_details']
        
        # 서브플롯 생성
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
        
        # 1. 원본 지표 값 막대 차트
        metric_names = list(metrics.keys())
        metric_values = list(metrics.values())
        
        bars1 = ax1.bar(metric_names, metric_values, color='skyblue', alpha=0.7)
        ax1.set_title('원본 지표 값', fontsize=14, fontweight='bold')
        ax1.set_ylabel('값')
        ax1.tick_params(axis='x', rotation=45)
        
        # 막대 위에 값 표시
        for bar, value in zip(bars1, metric_values):
            ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                    f'{value:.3f}', ha='center', va='bottom')
        
        # 2. 정규화된 점수 막대 차트
        normalized_scores = [details['normalized_score'] for details in score_details.values()]
        
        bars2 = ax2.bar(metric_names, normalized_scores, color='lightgreen', alpha=0.7)
        ax2.set_title('정규화된 점수 (0-100)', fontsize=14, fontweight='bold')
        ax2.set_ylabel('점수')
        ax2.set_ylim(0, 100)
        ax2.tick_params(axis='x', rotation=45)
        
        # 막대 위에 값 표시
        for bar, score in zip(bars2, normalized_scores):
            ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 2,
                    f'{score:.1f}', ha='center', va='bottom')
        
        # 3. 가중치 적용 점수 막대 차트
        weighted_scores = [details['weighted_score'] for details in score_details.values()]
        weights = [details['weight'] for details in score_details.values()]
        
        bars3 = ax3.bar(metric_names, weighted_scores, color='orange', alpha=0.7)
        ax3.set_title('가중치 적용 점수', fontsize=14, fontweight='bold')
        ax3.set_ylabel('가중 점수')
        ax3.tick_params(axis='x', rotation=45)
        
        # 막대 위에 가중치 표시
        for bar, score, weight in zip(bars3, weighted_scores, weights):
            ax3.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.5,
                    f'{score:.1f}\n(w:{weight})', ha='center', va='bottom', fontsize=9)
        
        # 4. 종합 점수 파이 차트
        comprehensive_score = results['comprehensive_score']
        grade = results['grade']
        
        # 등급별 색상 설정
        grade_colors = {'A': '#2e8b57', 'B': '#4169e1', 'C': '#ffa500', 'D': '#ff6347', 'F': '#dc143c'}
        
        ax4.pie([comprehensive_score, 100 - comprehensive_score], 
                labels=[f'{grade}등급\n{comprehensive_score:.1f}점', ''],
                colors=[grade_colors.get(grade, '#666666'), '#f0f0f0'],
                autopct='%1.1f%%', startangle=90)
        ax4.set_title('종합 점수', fontsize=14, fontweight='bold')
        
        plt.tight_layout()
        
        # 저장
        if save_path is None:
            save_path = os.path.join(self.images_dir, f"metrics_comparison_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png")
        
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
        
        print(f"지표 비교 차트가 저장되었습니다: {save_path}")
    
    def visualize_3d_model(self, model: Dict, title: str = "3D Model", 
                          save_path: Optional[str] = None) -> None:
        """
        3D 모델을 시각화합니다.
        
        Args:
            model (Dict): 3D 모델 정보
            title (str): 시각화 제목
            save_path (Optional[str]): 저장 경로
        """
        try:
            vertices = model['vertices']
            faces = model['faces']
            
            if len(vertices) == 0:
                print("시각화할 정점이 없습니다.")
                return
            
            # trimesh를 사용한 3D 시각화
            mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
            
            # 3D 플롯 생성
            fig = plt.figure(figsize=(12, 8))
            ax = fig.add_subplot(111, projection='3d')
            
            # 메시 플롯
            ax.plot_trisurf(vertices[:, 0], vertices[:, 1], vertices[:, 2], 
                           triangles=faces, alpha=0.8, cmap='viridis')
            
            # 바운딩 박스 표시
            bbox = model.get('bounding_box', None)
            if bbox is not None:
                self._plot_bounding_box(ax, bbox)
            
            ax.set_title(title, fontsize=14, fontweight='bold')
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            ax.set_zlabel('Z')
            
            # 축 비율 동일하게 설정
            self._set_equal_axes(ax, vertices)
            
            # 저장
            if save_path is None:
                save_path = os.path.join(self.images_dir, f"3d_model_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png")
            
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            plt.show()
            
            print(f"3D 모델 시각화가 저장되었습니다: {save_path}")
            
        except Exception as e:
            print(f"3D 모델 시각화 중 오류 발생: {e}")
    
    def create_performance_dashboard(self, results: Dict) -> str:
        """
        성능 대시보드를 생성합니다.
        
        Args:
            results (Dict): 평가 결과
            
        Returns:
            str: 대시보드 HTML 파일 경로
        """
        # Plotly 대시보드 생성
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=('종합 점수', '지표별 점수', '가중치 분포', '성능 등급'),
            specs=[[{"type": "indicator"}, {"type": "bar"}],
                   [{"type": "pie"}, {"type": "indicator"}]]
        )
        
        # 1. 종합 점수 게이지
        comprehensive_score = results['comprehensive_score']
        fig.add_trace(
            go.Indicator(
                mode="gauge+number+delta",
                value=comprehensive_score,
                domain={'x': [0, 1], 'y': [0, 1]},
                title={'text': "종합 점수"},
                gauge={
                    'axis': {'range': [None, 100]},
                    'bar': {'color': "darkblue"},
                    'steps': [
                        {'range': [0, 60], 'color': "lightgray"},
                        {'range': [60, 80], 'color': "yellow"},
                        {'range': [80, 100], 'color': "green"}
                    ],
                    'threshold': {
                        'line': {'color': "red", 'width': 4},
                        'thickness': 0.75,
                        'value': 90
                    }
                }
            ),
            row=1, col=1
        )
        
        # 2. 지표별 점수 막대 차트
        score_details = results['score_details']
        metric_names = list(score_details.keys())
        normalized_scores = [details['normalized_score'] for details in score_details.values()]
        
        fig.add_trace(
            go.Bar(
                x=metric_names,
                y=normalized_scores,
                name="정규화된 점수",
                marker_color='lightblue'
            ),
            row=1, col=2
        )
        
        # 3. 가중치 분포 파이 차트
        weights = [details['weight'] for details in score_details.values()]
        
        fig.add_trace(
            go.Pie(
                labels=metric_names,
                values=weights,
                name="가중치 분포"
            ),
            row=2, col=1
        )
        
        # 4. 성능 등급 표시
        grade = results['grade']
        grade_colors = {'A': '#2e8b57', 'B': '#4169e1', 'C': '#ffa500', 'D': '#ff6347', 'F': '#dc143c'}
        
        fig.add_trace(
            go.Indicator(
                mode="number",
                value=0,  # 등급은 텍스트로 표시
                title={'text': f"성능 등급: {grade}"},
                number={'font': {'size': 50, 'color': grade_colors.get(grade, '#666666')}}
            ),
            row=2, col=2
        )
        
        # 레이아웃 업데이트
        fig.update_layout(
            title_text="3D 객체인식 평가 대시보드",
            title_x=0.5,
            height=800,
            showlegend=False
        )
        
        # HTML 파일로 저장
        html_file = os.path.join(self.reports_dir, f"dashboard_{datetime.now().strftime('%Y%m%d_%H%M%S')}.html")
        fig.write_html(html_file)
        
        print(f"성능 대시보드가 생성되었습니다: {html_file}")
        
        return html_file
    
    def generate_html_report(self, results: Dict) -> str:
        """
        HTML 형식의 상세 리포트를 생성합니다.
        
        Args:
            results (Dict): 평가 결과
            
        Returns:
            str: HTML 리포트 파일 경로
        """
        html_content = f"""
        <!DOCTYPE html>
        <html lang="ko">
        <head>
            <meta charset="UTF-8">
            <meta name="viewport" content="width=device-width, initial-scale=1.0">
            <title>3D 객체인식 평가 결과</title>
            <style>
                body {{
                    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
                    margin: 0;
                    padding: 20px;
                    background-color: #f5f5f5;
                }}
                .container {{
                    max-width: 1200px;
                    margin: 0 auto;
                    background-color: white;
                    padding: 30px;
                    border-radius: 10px;
                    box-shadow: 0 0 20px rgba(0,0,0,0.1);
                }}
                .header {{
                    text-align: center;
                    margin-bottom: 30px;
                    padding: 20px;
                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                    color: white;
                    border-radius: 10px;
                }}
                .score-section {{
                    display: flex;
                    justify-content: space-around;
                    margin: 30px 0;
                    padding: 20px;
                    background-color: #f8f9fa;
                    border-radius: 10px;
                }}
                .score-item {{
                    text-align: center;
                }}
                .score-value {{
                    font-size: 2.5em;
                    font-weight: bold;
                    color: #2e8b57;
                }}
                .grade-value {{
                    font-size: 2em;
                    font-weight: bold;
                    color: #4169e1;
                }}
                .metrics-grid {{
                    display: grid;
                    grid-template-columns: repeat(auto-fit, minmax(300px, 1fr));
                    gap: 20px;
                    margin: 30px 0;
                }}
                .metric-card {{
                    background-color: #f8f9fa;
                    padding: 20px;
                    border-radius: 10px;
                    border-left: 5px solid #007bff;
                }}
                .metric-name {{
                    font-weight: bold;
                    font-size: 1.2em;
                    color: #333;
                }}
                .metric-value {{
                    font-size: 1.5em;
                    color: #666;
                    margin: 10px 0;
                }}
                .metric-details {{
                    font-size: 0.9em;
                    color: #888;
                }}
                .summary-table {{
                    width: 100%;
                    border-collapse: collapse;
                    margin: 30px 0;
                }}
                .summary-table th, .summary-table td {{
                    border: 1px solid #ddd;
                    padding: 12px;
                    text-align: left;
                }}
                .summary-table th {{
                    background-color: #f2f2f2;
                    font-weight: bold;
                }}
                .footer {{
                    text-align: center;
                    margin-top: 30px;
                    padding: 20px;
                    background-color: #f8f9fa;
                    border-radius: 10px;
                    color: #666;
                }}
            </style>
        </head>
        <body>
            <div class="container">
                <div class="header">
                    <h1>3D 객체인식 평가 결과</h1>
                    <p>평가 시간: {results['evaluation_timestamp']}</p>
                    <p>모델 파일: {os.path.basename(results['model_path'])}</p>
                    <p>참조 파일: {os.path.basename(results['reference_path'])}</p>
                </div>
                
                <div class="score-section">
                    <div class="score-item">
                        <div class="score-value">{results['comprehensive_score']:.1f}점</div>
                        <div>종합 점수</div>
                    </div>
                    <div class="score-item">
                        <div class="grade-value">{results['grade']}등급</div>
                        <div>성능 등급</div>
                    </div>
                </div>
                
                <div class="metrics-grid">
                    {self._generate_metric_cards_html(results['score_details'])}
                </div>
                
                <h2>상세 결과</h2>
                <table class="summary-table">
                    <thead>
                        <tr>
                            <th>지표</th>
                            <th>원본 값</th>
                            <th>정규화 점수</th>
                            <th>가중치</th>
                            <th>가중 점수</th>
                        </tr>
                    </thead>
                    <tbody>
                        {self._generate_summary_table_html(results['score_details'])}
                    </tbody>
                </table>
                
                <div class="footer">
                    <p>이 리포트는 3D 객체인식 평가 시스템에 의해 자동 생성되었습니다.</p>
                    <p>생성 시간: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
                </div>
            </div>
        </body>
        </html>
        """
        
        # HTML 파일로 저장
        html_file = os.path.join(self.reports_dir, f"evaluation_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.html")
        
        with open(html_file, 'w', encoding='utf-8') as f:
            f.write(html_content)
        
        print(f"HTML 리포트가 생성되었습니다: {html_file}")
        
        return html_file
    
    def _plot_bounding_box(self, ax, bbox):
        """바운딩 박스를 플롯합니다."""
        min_coords, max_coords = bbox
        
        # 바운딩 박스의 8개 꼭짓점
        vertices = np.array([
            [min_coords[0], min_coords[1], min_coords[2]],
            [max_coords[0], min_coords[1], min_coords[2]],
            [max_coords[0], max_coords[1], min_coords[2]],
            [min_coords[0], max_coords[1], min_coords[2]],
            [min_coords[0], min_coords[1], max_coords[2]],
            [max_coords[0], min_coords[1], max_coords[2]],
            [max_coords[0], max_coords[1], max_coords[2]],
            [min_coords[0], max_coords[1], max_coords[2]]
        ])
        
        # 바운딩 박스의 12개 모서리
        edges = [
            [0, 1], [1, 2], [2, 3], [3, 0],  # 하단
            [4, 5], [5, 6], [6, 7], [7, 4],  # 상단
            [0, 4], [1, 5], [2, 6], [3, 7]   # 수직
        ]
        
        # 모서리 그리기
        for edge in edges:
            points = vertices[edge]
            ax.plot3D(*points.T, color='red', linewidth=2, alpha=0.8)
    
    def _set_equal_axes(self, ax, vertices):
        """축 비율을 동일하게 설정합니다."""
        max_range = np.array([vertices[:, 0].max() - vertices[:, 0].min(),
                             vertices[:, 1].max() - vertices[:, 1].min(),
                             vertices[:, 2].max() - vertices[:, 2].min()]).max() / 2.0
        
        mid_x = (vertices[:, 0].max() + vertices[:, 0].min()) * 0.5
        mid_y = (vertices[:, 1].max() + vertices[:, 1].min()) * 0.5
        mid_z = (vertices[:, 2].max() + vertices[:, 2].min()) * 0.5
        
        ax.set_xlim(mid_x - max_range, mid_x + max_range)
        ax.set_ylim(mid_y - max_range, mid_y + max_range)
        ax.set_zlim(mid_z - max_range, mid_z + max_range)
    
    def _generate_metric_cards_html(self, score_details: Dict) -> str:
        """지표 카드 HTML을 생성합니다."""
        html = ""
        for metric_name, details in score_details.items():
            html += f"""
            <div class="metric-card">
                <div class="metric-name">{metric_name}</div>
                <div class="metric-value">{details['normalized_score']:.1f}점</div>
                <div class="metric-details">
                    원본값: {details['raw_value']:.4f}<br>
                    가중치: {details['weight']}<br>
                    임계값: {details['threshold']}
                </div>
            </div>
            """
        return html
    
    def _generate_summary_table_html(self, score_details: Dict) -> str:
        """요약 테이블 HTML을 생성합니다."""
        html = ""
        for metric_name, details in score_details.items():
            html += f"""
            <tr>
                <td>{metric_name}</td>
                <td>{details['raw_value']:.4f}</td>
                <td>{details['normalized_score']:.1f}</td>
                <td>{details['weight']}</td>
                <td>{details['weighted_score']:.2f}</td>
            </tr>
            """
        return html
    
    def create_comparison_chart(self, results_list: List[Dict], save_path: Optional[str] = None) -> None:
        """
        여러 모델의 평가 결과를 비교하는 차트를 생성합니다.
        
        Args:
            results_list (List[Dict]): 여러 평가 결과 리스트
            save_path (Optional[str]): 저장 경로
        """
        if not results_list:
            print("비교할 결과가 없습니다.")
            return
        
        # 데이터 준비
        model_names = [f"Model {i+1}" for i in range(len(results_list))]
        comprehensive_scores = [results['comprehensive_score'] for results in results_list]
        grades = [results['grade'] for results in results_list]
        
        # 서브플롯 생성
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # 1. 종합 점수 비교
        bars = ax1.bar(model_names, comprehensive_scores, color='lightblue', alpha=0.7)
        ax1.set_title('모델별 종합 점수 비교', fontsize=14, fontweight='bold')
        ax1.set_ylabel('종합 점수')
        ax1.set_ylim(0, 100)
        
        # 막대 위에 점수 표시
        for bar, score in zip(bars, comprehensive_scores):
            ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1,
                    f'{score:.1f}', ha='center', va='bottom')
        
        # 2. 등급 분포
        grade_counts = {}
        for grade in grades:
            grade_counts[grade] = grade_counts.get(grade, 0) + 1
        
        ax2.pie(grade_counts.values(), labels=grade_counts.keys(), autopct='%1.1f%%')
        ax2.set_title('성능 등급 분포', fontsize=14, fontweight='bold')
        
        plt.tight_layout()
        
        # 저장
        if save_path is None:
            save_path = os.path.join(self.images_dir, f"model_comparison_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png")
        
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
        
        print(f"모델 비교 차트가 저장되었습니다: {save_path}")
