import matplotlib.pyplot as plt
import plotly.graph_objects as go
import plotly.express as px
import pandas as pd
import numpy as np
import streamlit as st
from config import SOIL_TYPES, STRENGTH_PARAMETERS

class SoilProfileVisualizer:
    def __init__(self):
        self.soil_colors = {
            "soft clay": "#8B4513",
            "medium clay": "#A0522D", 
            "stiff clay": "#D2691E",
            "very stiff clay": "#CD853F",
            "hard clay": "#DEB887",
            "loose sand": "#F4A460",
            "medium dense sand": "#DAA520",
            "dense sand": "#B8860B",
            "very dense sand": "#CD853F",
            "soft silt": "#DDA0DD",
            "medium silt": "#BA55D3",
            "stiff silt": "#9370DB",
            "loose gravel": "#696969",
            "dense gravel": "#2F4F4F",
            "weathered rock": "#708090",
            "soft rock": "#2F4F4F",
            "hard rock": "#36454F"
        }
    
    def create_soil_profile_plot(self, soil_data):
        """Create interactive soil profile visualization"""
        if not soil_data or "soil_layers" not in soil_data:
            return None
        
        layers = soil_data["soil_layers"]
        
        fig = go.Figure()
        
        # Add soil layers
        for i, layer in enumerate(layers):
            depth_from = layer.get("depth_from", 0)
            depth_to = layer.get("depth_to", 0)
            soil_type = layer.get("soil_type", "unknown")
            description = layer.get("description", "")
            strength_value = layer.get("strength_value", "N/A")
            strength_param = layer.get("strength_parameter", "")
            
            # Get color
            color = self.soil_colors.get(soil_type.lower(), "#CCCCCC")
            
            # Create layer rectangle
            fig.add_shape(
                type="rect",
                x0=0, x1=1,
                y0=-depth_to, y1=-depth_from,
                fillcolor=color,
                line=dict(color="black", width=1),
                opacity=0.8
            )
            
            # Add layer text with enhanced parameters
            mid_depth = -(depth_from + depth_to) / 2
            
            # Build text with available parameters
            text_lines = [f"{layer.get('consistency', '')} {soil_type}".strip()]
            
            # Add strength parameters
            if strength_param and strength_value is not None:
                text_lines.append(f"{strength_param}: {strength_value}")
            
            # Add calculated Su if available
            if layer.get("calculated_su"):
                text_lines.append(f"Su: {layer['calculated_su']:.0f} kPa*")
            
            # Add friction angle if available
            if layer.get("friction_angle"):
                text_lines.append(f"φ: {layer['friction_angle']:.1f}°*")
            
            fig.add_annotation(
                x=0.5, y=mid_depth,
                text="<br>".join(text_lines),
                showarrow=False,
                font=dict(size=9, color="white"),
                bgcolor="rgba(0,0,0,0.6)",
                bordercolor="white",
                borderwidth=1
            )
        
        # Add depth markers
        max_depth = max([layer.get("depth_to", 0) for layer in layers])
        depth_ticks = list(range(0, int(max_depth) + 5, 5))
        
        fig.update_layout(
            title="Soil Profile",
            xaxis=dict(
                range=[0, 1],
                showticklabels=False,
                showgrid=False,
                zeroline=False
            ),
            yaxis=dict(
                title="Depth (m)",
                range=[-max_depth - 2, 2],
                tickvals=[-d for d in depth_ticks],
                ticktext=[str(d) for d in depth_ticks],
                showgrid=True,
                gridcolor="lightgray"
            ),
            width=400,
            height=600,
            margin=dict(l=50, r=50, t=50, b=50)
        )
        
        # Add water table if present
        if "water_table" in soil_data and soil_data["water_table"].get("depth"):
            wt_depth = soil_data["water_table"]["depth"]
            fig.add_hline(
                y=-wt_depth,
                line_dash="dash",
                line_color="blue",
                annotation_text="Water Table",
                annotation_position="right"
            )
        
        return fig
    
    def create_strength_profile_plot(self, soil_data):
        """Create strength parameter vs depth plot"""
        if not soil_data or "soil_layers" not in soil_data:
            return None
        
        layers = soil_data["soil_layers"]
        
        depths = []
        strengths = []
        soil_types = []
        
        for layer in layers:
            depth_from = layer.get("depth_from", 0)
            depth_to = layer.get("depth_to", 0)
            strength_value = layer.get("strength_value")
            soil_type = layer.get("soil_type", "")
            
            if strength_value is not None:
                mid_depth = (depth_from + depth_to) / 2
                depths.append(mid_depth)
                strengths.append(strength_value)
                soil_types.append(soil_type)
        
        if not depths:
            return None
        
        fig = go.Figure()
        
        # Group by parameter type
        clay_depths = []
        clay_strengths = []
        sand_depths = []
        sand_strengths = []
        
        for i, soil_type in enumerate(soil_types):
            if "clay" in soil_type.lower():
                clay_depths.append(depths[i])
                clay_strengths.append(strengths[i])
            else:
                sand_depths.append(depths[i])
                sand_strengths.append(strengths[i])
        
        # Add traces
        if clay_depths:
            # Create custom hover text for Su values
            clay_hover_text = [f"Depth: {d:.1f}m<br>Su: {s:.1f} kPa" for d, s in zip(clay_depths, clay_strengths)]
            
            fig.add_trace(go.Scatter(
                x=clay_strengths,
                y=clay_depths,
                mode='markers+lines',
                name='Su (kPa)',
                marker=dict(color='brown', size=8),
                line=dict(color='brown'),
                hovertemplate='%{customdata}<extra></extra>',
                customdata=clay_hover_text
            ))
        
        if sand_depths:
            # Create custom hover text for SPT-N values
            sand_hover_text = [f"Depth: {d:.1f}m<br>SPT-N: {s:.0f} blows/30cm" for d, s in zip(sand_depths, sand_strengths)]
            
            fig.add_trace(go.Scatter(
                x=sand_strengths,
                y=sand_depths,
                mode='markers+lines',
                name='SPT-N (blows/30cm)',
                marker=dict(color='gold', size=8),
                line=dict(color='gold'),
                hovertemplate='%{customdata}<extra></extra>',
                customdata=sand_hover_text
            ))
        
        # Determine primary axis title based on data
        if clay_depths and sand_depths:
            xaxis_title = "Strength Value (Su in kPa / SPT-N)"
        elif clay_depths:
            xaxis_title = "Undrained Shear Strength, Su (kPa)"
        elif sand_depths:
            xaxis_title = "SPT-N Value (blows/30cm)"
        else:
            xaxis_title = "Strength Value"
        
        fig.update_layout(
            title="Strength Parameters vs Depth",
            xaxis_title=xaxis_title,
            yaxis_title="Depth (m)",
            yaxis=dict(autorange='reversed'),
            width=500,
            height=600,
            showlegend=True,
            legend=dict(
                yanchor="top",
                y=0.99,
                xanchor="left",
                x=0.01
            )
        )
        
        return fig
    
    def create_layer_summary_table(self, soil_data):
        """Create summary table of soil layers"""
        if not soil_data or "soil_layers" not in soil_data:
            return None
        
        layers = soil_data["soil_layers"]
        
        df_data = []
        for layer in layers:
            # Build strength info with units
            strength_info = ""
            if layer.get("strength_parameter") and layer.get("strength_value") is not None:
                param = layer['strength_parameter']
                value = layer['strength_value']
                
                # Add units based on parameter type
                if param == "Su":
                    strength_info = f"Su: {value:.1f} kPa"
                elif param == "SPT-N":
                    strength_info = f"SPT-N: {value:.0f} blows/30cm"
                else:
                    strength_info = f"{param}: {value}"
            
            # Add calculated parameters
            calc_params = []
            if layer.get("calculated_su"):
                calc_params.append(f"Su: {layer['calculated_su']:.0f} kPa (calc)")
            if layer.get("friction_angle"):
                calc_params.append(f"φ: {layer['friction_angle']:.1f}° (calc)")
            
            if calc_params:
                strength_info += f" | {' | '.join(calc_params)}"
            
            df_data.append({
                "Layer": layer.get("layer_id", ""),
                "Depth From (m)": layer.get("depth_from", ""),
                "Depth To (m)": layer.get("depth_to", ""),
                "Soil Type": f"{layer.get('consistency', '')} {layer.get('soil_type', '')}".strip(),
                "Description": layer.get("description", ""),
                "Strength Parameters": strength_info,
                "Color": layer.get("color", ""),
                "Moisture": layer.get("moisture", ""),
                "Notes": layer.get("su_source", "") or layer.get("friction_angle_source", "") or ""
            })
        
        return pd.DataFrame(df_data)
    
    def export_profile_data(self, soil_data, format="csv"):
        """Export soil profile data"""
        df = self.create_layer_summary_table(soil_data)
        
        if format == "csv":
            return df.to_csv(index=False)
        elif format == "json":
            return df.to_json(orient="records", indent=2)
        else:
            return df.to_string(index=False)