import gradio as gr
import matplotlib.pyplot as plt
from datetime import datetime
from dateutil import parser
from io import BytesIO
from PIL import Image
from geopy.geocoders import Nominatim
from timezonefinder import TimezoneFinder
import pytz
import swisseph as swe

# Initialize Swiss Ephemeris
swe.set_ephe_path(None)

# Russian translations for planets
PLANET_RU = {
    'Sun': 'Солнце', 'Moon': 'Луна', 'Mercury': 'Меркурий',
    'Venus': 'Венера', 'Mars': 'Марс', 
    'Jupiter': 'Юпитер', 'Saturn': 'Сатурн'
}

# Planet symbols for plotting
PLANET_SYMBOLS = {
    'Sun': '☉', 'Moon': '☾', 'Mercury': '☿', 'Venus': '♀',
    'Mars': '♂', 'Jupiter': '♃', 'Saturn': '♄'
}

# Zodiac signs in Russian
ZODIAC_SIGNS = [
    "Овен", "Телец", "Близнецы", "Рак", "Лев", "Дева",
    "Весы", "Скорпион", "Стрелец", "Козерог", "Водолей", "Рыбы"
]

def parse_query(query):
    """Parse the query into date, time, and location."""
    if not query.startswith("PLadder "):
        return None, None, "Query must start with 'PLadder'"
    
    try:
        parts = query.split(maxsplit=3)
        if len(parts) < 4:
            return None, None, "Incomplete query (need date, time, and location)"
        
        _, date_str, time_str, location = parts
        dt = parser.parse(f"{date_str} {time_str}")
        return dt, location, None
    except ValueError as e:
        return None, None, f"Invalid format: {str(e)}"

def get_utc_time(dt, location):
    """Convert local time to UTC using location's time zone."""
    geolocator = Nominatim(user_agent="pladder_app")
    try:
        loc = geolocator.geocode(location, timeout=10)
        if not loc:
            return None, None, None, "Location not found"
        
        lat, lon = loc.latitude, loc.longitude
        tz_str = TimezoneFinder().timezone_at(lng=lon, lat=lat)
        if not tz_str:
            return None, None, None, "Time zone not found"
        
        tz = pytz.timezone(tz_str)
        local_dt = tz.localize(dt)
        utc_dt = local_dt.astimezone(pytz.UTC)
        return utc_dt, lat, lon, None
    except Exception as e:
        return None, None, None, f"Error: {str(e)}"

def format_coords(lat, lon):
    """Format coordinates as degrees, minutes, seconds."""
    def dms(value, pos_dir, neg_dir):
        direction = pos_dir if value >= 0 else neg_dir
        abs_value = abs(value)
        degrees = int(abs_value)
        minutes = int((abs_value - degrees) * 60)
        seconds = int(round(((abs_value - degrees) * 60 - minutes) * 60))
        
        # Handle rounding overflow
        if seconds >= 60:
            seconds -= 60
            minutes += 1
        if minutes >= 60:
            minutes -= 60
            degrees += 1
            
        return f"{degrees}°{minutes:02}'{seconds:02}\" {direction}"
    
    return f"{dms(lat, 'N', 'S')}, {dms(lon, 'E', 'W')}"

def lon_to_sign(lon_deg):
    """
    Convert ecliptic longitude to zodiac sign with position.
    Now includes seconds in the output.
    """
    sign_idx = int(lon_deg // 30)
    degrees_in_sign = lon_deg % 30
    degrees = int(degrees_in_sign)
    minutes = int((degrees_in_sign - degrees) * 60)
    seconds = int(round(((degrees_in_sign - degrees) * 60 - minutes) * 60))
    
    # Handle rounding overflow
    if seconds >= 60:
        seconds -= 60
        minutes += 1
    if minutes >= 60:
        minutes -= 60
        degrees += 1
    
    return f"{ZODIAC_SIGNS[sign_idx]} {degrees}°{minutes:02}'{seconds:02}\""

def PLadder_ZSizes(utc_dt, lat, lon):
    """Calculate Planetary Ladder and Zone Sizes using Swiss Ephemeris."""
    if not -13000 <= utc_dt.year <= 17000:
        return {"error": "Date out of supported range (-13,000–17,000 CE)"}
    
    # Planet mapping with Swiss Ephemeris constants
    PLANET_OBJECTS = {
        'Sun': swe.SUN, 'Moon': swe.MOON, 'Mercury': swe.MERCURY,
        'Venus': swe
.VENUS, 'Mars': swe.MARS, 
        'Jupiter': swe.JUPITER, 'Saturn': swe.SATURN
    }
    
    # Calculate Julian Day
    jd_utc = swe.julday(
        utc_dt.year, utc_dt.month, utc_dt.day,
        utc_dt.hour + utc_dt.minute/60 + utc_dt.second/3600
    )
    
    # Calculate planetary positions
    longitudes = {}
    for planet, planet_id in PLANET_OBJECTS.items():
        flags = swe.FLG_SWIEPH | swe.FLG_SPEED
        xx, _ = swe.calc_ut(jd_utc, planet_id, flags)
        lon = xx[0] % 360  # Normalize to 0-360°
        longitudes[planet] = lon
    
    # Sort planets by longitude
    sorted_planets = sorted(longitudes.items(), key=lambda x: x[1])
    PLadder = [p for p, _ in sorted_planets]
    sorted_lons = [lon for _, lon in sorted_planets]
    
    # Calculate zone sizes
    zone_sizes = [sorted_lons[0]]  # First zone
    zone_sizes.extend(sorted_lons[i+1] - sorted_lons[i] for i in range(6))  # Middle zones
    zone_sizes.append(360 - sorted_lons[-1])  # Last zone
    
    # Classify zone sizes
    ZSizes = []
    for i, size in enumerate(zone_sizes):
        # Get bordering planets
        if i == 0:
            bord = [PLadder[0]]
        elif i == len(zone_sizes)-1:
            bord = [PLadder[-1]]
        else:
            bord = [PLadder[i-1], PLadder[i]]
        
        # Determine zone classification
        if any(p in ['Sun', 'Moon'] for p in bord):
            X = 7
        elif any(p in ['Mercury', 'Venus', 'Mars'] for p in bord):
            X = 6
        else:
            X = 5
            
        # Format size with seconds
        d = int(size)
        m = int((size - d) * 60)
        s = int(round(((size - d) * 60 - m) * 60))
        
        # Handle rounding overflow
        if s >= 60:
            s -= 60
            m += 1
        if m >= 60:
            m -= 60
            d += 1
            
        classification = (
            'Swallowed' if size <= 1 else
            'Tiny' if size <= X else
            'Small' if size <= 40 else
            'Ideal' if 50 <= size <= 52 else
            'Normal' if size < 60 else
            'Big'
        )
        
        ZSizes.append((f"{d}°{m:02}'{s:02}\"", classification))
    
    return {
        'PLadder': PLadder,
        'ZSizes': ZSizes,
        'longitudes': longitudes  # Raw degrees for calculations
    }

def plot_pladder(PLadder):
    """Generate the original version of the planetary ladder visualization."""
    fig, ax = plt.subplots(figsize=(6, 6))
    
    # Draw the main triangle
    ax.plot([0, 1.5, 3, 0], [0, 3, 0, 0], 'k-', linewidth=2)
    
    # Draw horizontal divisions (original style)
    ax.plot([0.5, 2.5], [1, 1], 'k--')
    ax.plot([1, 2], [2, 2], 'k--')
    
    # Original planet symbol positions
    symbol_positions = [
        (-0.2, 0.2), (0.3, 1.2), (0.8, 2.2),
        (1.5, 3.2), (2.2, 2.2), (2.7, 1.2), (3.2, 0.2)
    ]
    
    # Add planet symbols
    for (x, y), planet in zip(symbol_positions, PLadder):
        ax.text(x, y, PLANET_SYMBOLS[planet], 
               ha='center', va='center', 
               fontsize=24)
    
    # Configure plot appearance (original style)
    ax.set_xlim(-0.5, 3.5)
    ax.set_ylim(-0.5, 3.5)
    ax.set_aspect('equal')
    ax.axis('off')
    
    return fig

def chat_interface(query):
    """Process the user query and return text and plot."""
    # Parse input
    dt, location, error = parse_query(query)
    if error:
        return error, None
    
    # Get UTC time and coordinates
    utc_dt, lat, lon, error = get_utc_time(dt, location)
    if error:
        return error, None
    
    # Calculate planetary positions
    result = PLadder_ZSizes(utc_dt, lat, lon)
    if "error" in result:
        return result["error"], None
    
    # Format output
    PLadder = result["PLadder"]
    ZSizes = result["ZSizes"]
    longitudes = result["longitudes"]
    
    # Generate planet list text with full DMS
    planet_list = "\n".join(
        f"{PLANET_RU[p]}: {lon_to_sign(longitudes[p])}"
        for p in PLadder
    )
    
    # Generate zone sizes text
    zones_text = "\n".join(
        f"Zone {i+1}: {size} ({cls})"
        for i, (size, cls) in enumerate(ZSizes)
    )
    
    # Generate coordinates text
    coords_text = format_coords(lat, lon)
    
    # Create visualization (original style)
    fig = plot_pladder(PLadder)
    buf = BytesIO()
    fig.savefig(buf, format='png', dpi=120, bbox_inches='tight')
    buf.seek(0)
    img = Image.open(buf)
    plt.close(fig)
    
    # Compose final output
    output_text = (
        f"Planetary Ladder:\n{planet_list}\n\n"
        f"Zone Sizes:\n{zones_text}\n\n"
        f"Coordinates: {coords_text}\n"
        f"Calculation Time: {utc_dt.strftime('%Y-%m-%d %H:%M:%S UTC')}"
    )
    
    return output_text, img

# Gradio Interface
with gr.Blocks(title="Planetary Ladder Calculator") as interface:
    gr.Markdown("## Planetary Ladder Calculator")
    with gr.Row():
        with gr.Column(scale=2):
            output_text = gr.Textbox(label="Astrological Data", lines=12)
        with gr.Column(scale=1):
            output_image = gr.Image(label="Visualization")
    with gr.Row():
        query_text = gr.Textbox(
            label="Input Query",
            placeholder="Example for Elon Mask: PLadder 28-06-1971 12:03:00 Pretoria",
            max_lines=1
        )
    with gr.Row():
        submit_button = gr.Button("Calculate", variant="primary")
    
    submit_button.click(
        fn=chat_interface,
        inputs=query_text,
        outputs=[output_text, output_image]
    )

if __name__ == "__main__":
    interface.launch()