import gradio as gr
import swisseph as swe
from datetime import datetime, timedelta
from geopy.geocoders import Nominatim
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import numpy as np

# Initialize Swiss Ephemeris
swe.set_ephe_path('')

def get_coordinates(location_name):
    """Get latitude and longitude from location name."""
    try:
        geolocator = Nominatim(user_agent="rashi_chart_generator")
        location = geolocator.geocode(location_name)
        if location:
            return location.latitude, location.longitude, location.address
        else:
            return 28.6139, 77.2090, "New Delhi, India"
    except Exception as e:
        return 28.6139, 77.2090, "New Delhi, India"

def get_julian_day(date_str, lat, lon):
    """Calculate Julian Day for noon (12:00 PM)."""
    try:
        date_obj = datetime.strptime(date_str, "%Y-%m-%d")
        year = date_obj.year
        month = date_obj.month
        day = date_obj.day
        hour = 12.0
        jd = swe.julday(year, month, day, hour)
        return jd
    except Exception as e:
        print(f"Error calculating Julian Day: {e}")
        return None

def calculate_planet_positions(jd):
    """Calculate positions of all planets."""
    planets = {
        'Sun': swe.SUN,
        'Moon': swe.MOON,
        'Mars': swe.MARS,
        'Mercury': swe.MERCURY,
        'Jupiter': swe.JUPITER,
        'Venus': swe.VENUS,
        'Saturn': swe.SATURN,
        'Rahu': swe.MEAN_NODE,
    }
    positions = {}
    for planet_name, planet_id in planets.items():
        try:
            result = swe.calc_ut(jd, planet_id)
            longitude = result[0][0]
            positions[planet_name] = longitude
        except Exception as e:
            positions[planet_name] = 0.0
    positions['Ketu'] = (positions['Rahu'] + 180) % 360
    return positions

def calculate_moon_sign(birth_date, birth_location):
    """Calculate Moon sign (Rashi) from birth date."""
    lat, lon, address = get_coordinates(birth_location)
    jd = get_julian_day(birth_date, lat, lon)
    if jd is None:
        return None, None
    result = swe.calc_ut(jd, swe.MOON)
    moon_longitude = result[0][0]
    moon_rashi = int(moon_longitude / 30) + 1
    if moon_rashi > 12:
        moon_rashi = 12
    return moon_rashi, moon_longitude

def degree_to_rashi(degree):
    """Convert degree to Rashi number (1-12)."""
    rashi_num = int(degree / 30) + 1
    if rashi_num > 12:
        rashi_num = 12
    return rashi_num

def get_rashi_name(rashi_num):
    """Get Rashi name."""
    rashi_names = {
        1: 'Mesha', 2: 'Vrishabha', 3: 'Mithuna', 4: 'Karka',
        5: 'Simha', 6: 'Kanya', 7: 'Tula', 8: 'Vrishchika',
        9: 'Dhanu', 10: 'Makara', 11: 'Kumbha', 12: 'Meena'
    }
    return rashi_names.get(rashi_num, 'Unknown')

def organize_planets_by_house(positions, moon_rashi):
    """Organize current planets into houses relative to birth Moon sign."""
    house_planets = {i: [] for i in range(1, 13)}
    for planet, degree in positions.items():
        planet_rashi = degree_to_rashi(degree)
        house = ((planet_rashi - moon_rashi) % 12) + 1
        if house > 12:
            house = house - 12
        house_planets[house].append(planet)
    return house_planets

def get_short_planet_names(house_planets):
    """Convert to abbreviations."""
    short_names = {
        'Sun': 'Su', 'Moon': 'Mo', 'Mars': 'Ma', 'Mercury': 'Me',
        'Jupiter': 'Ju', 'Venus': 'Ve', 'Saturn': 'Sa',
        'Rahu': 'Ra', 'Ketu': 'Ke',
    }
    short_house_planets = {}
    for house, planets in house_planets.items():
        short_house_planets[house] = [short_names.get(p, p) for p in planets]
    return short_house_planets

def get_house_rashis(moon_rashi):
    """Get Rashi name for each house based on Moon sign."""
    house_rashis = {}
    for house in range(1, 13):
        rashi_num = ((moon_rashi + house - 2) % 12) + 1
        house_rashis[house] = get_rashi_name(rashi_num)
    return house_rashis

def draw_rashi_transit_chart(house_planets, house_rashis, birth_date, current_date, location, moon_rashi, save_path):
    """Draw square grid Rashi chart with current transits."""
    fig, ax = plt.subplots(1, 1, figsize=(12, 12))
    ax.set_xlim(0, 4)
    ax.set_ylim(0, 4)
    ax.set_aspect('equal')
    ax.axis('off')
    fig.patch.set_facecolor('#f5f5dc')
    
    moon_sign_name = get_rashi_name(moon_rashi)
    title_text = f"Rashi Chart (Moon Sign: {moon_sign_name})\n"
    title_text += f"Birth Date: {birth_date} | Current Date: {current_date}\n"
    title_text += f"Location: {location}"
    ax.text(2, 4.4, title_text, ha='center', va='top',
            fontsize=11, fontweight='bold', color='#8B4513')
    
    grid_layout = {
        12: (3,0), 1:(3,1), 2:(3,2), 3:(3,3),
        11:(2,0), 4:(2,3),
        10:(1,0), 5:(1,3),
        9:(0,0), 8:(0,1), 7:(0,2), 6:(0,3),
    }
    
    for house_num in range(1, 13):
        row, col = grid_layout[house_num]
        x = col
        y = 3 - row
        
        if house_num == 1:
            facecolor = '#FFE4B5'
            edgecolor = '#FF6347'
            linewidth = 3
        else:
            facecolor = '#FFFACD'
            edgecolor = '#8B4513'
            linewidth = 2
        
        rect = Rectangle((x, y), 1, 1, linewidth=linewidth,
                        edgecolor=edgecolor, facecolor=facecolor)
        ax.add_patch(rect)
        
        house_label = f"{house_num}"
        if house_num == 1:
            house_label += " (Lg)"
        ax.text(x + 0.08, y + 0.92, house_label,
                fontsize=9, ha='left', va='top',
                color='gray', fontweight='bold')
        
        rashi_name = house_rashis.get(house_num, '')
        ax.text(x + 0.5, y + 0.82, rashi_name,
                fontsize=9, ha='center', va='top',
                fontweight='bold', color='#006400')
        
        planets = house_planets.get(house_num, [])
        if planets:
            planet_text = '\n'.join(planets)
            ax.text(x + 0.5, y + 0.45, planet_text,
                    fontsize=11, ha='center', va='center',
                    color='red', fontweight='bold')
    
    center_text = f"RASI\n({moon_sign_name})"
    ax.text(2, 2, center_text, fontsize=16, ha='center', va='center',
            fontweight='bold', color='#8B4513')
    center_rect = Rectangle((1, 1), 2, 2, linewidth=2,
                            edgecolor='#8B4513', facecolor='#FFFACD',
                            linestyle='--')
    ax.add_patch(center_rect)
    
    legend_text = "Current Planetary Positions (Transit)\n"
    legend_text += "Su=Sun, Mo=Moon, Ma=Mars, Me=Mercury, Ju=Jupiter\n"
    legend_text += "Ve=Venus, Sa=Saturn, Ra=Rahu, Ke=Ketu\n"
    legend_text += "House 1 (Lg) = Lagna (Moon Sign at birth)"
    ax.text(2, -0.4, legend_text, ha='center', va='top',
            fontsize=8, style='italic', color='#8B4513',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.3))
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight', facecolor='#f5f5dc')
    plt.close()
    return save_path

def generate_rashi_chart(birth_date, birth_location, transit_date=None):
    """Main function to generate Rashi chart."""
    if transit_date is None or transit_date == "":
        transit_date = datetime.now().strftime('%Y-%m-%d')
    
    moon_rashi, moon_longitude = calculate_moon_sign(birth_date, birth_location)
    if moon_rashi is None:
        return None, "Failed to calculate Moon sign. Please check your inputs."
    
    moon_sign_name = get_rashi_name(moon_rashi)
    lat, lon, full_address = get_coordinates(birth_location)
    jd = get_julian_day(transit_date, lat, lon)
    if jd is None:
        return None, "Failed to calculate Julian Day."
    
    current_positions = calculate_planet_positions(jd)
    house_planets = organize_planets_by_house(current_positions, moon_rashi)
    short_house_planets = get_short_planet_names(house_planets)
    house_rashis = get_house_rashis(moon_rashi)
    
    chart_filename = f"rashi_chart_{birth_date.replace('-', '_')}_{transit_date.replace('-', '_')}.png"
    chart_path = draw_rashi_transit_chart(short_house_planets, house_rashis, 
                                          birth_date, transit_date, full_address, 
                                          moon_rashi, chart_filename)
    
    report = f"""
## 🌙 RASHI CHART REPORT

### Birth Information
- **Moon Sign (Rashi):** {moon_sign_name} (Rashi {moon_rashi})
- **Moon Position:** {moon_longitude:.2f}°
- **Birth Date:** {birth_date}
- **Birth Location:** {full_address}

### Transit Information
- **Transit Date:** {transit_date}

### Current Planetary Positions

| Planet | Degree | Rashi | Sign |
|--------|--------|-------|------|
"""
    
    for planet, degree in current_positions.items():
        rashi = degree_to_rashi(degree)
        rashi_name = get_rashi_name(rashi)
        report += f"| {planet} | {degree:.2f}° | {rashi} | {rashi_name} |\n"
    
    report += "\n### House-wise Planet Distribution\n\n"
    
    for house in range(1, 13):
        planets = short_house_planets.get(house, [])
        rashi_name = house_rashis.get(house, '')
        planet_list = ', '.join(planets) if planets else 'Empty'
        house_label = f"House {house}"
        if house == 1:
            house_label += " (Lagna)"
        report += f"**{house_label}** ({rashi_name}): {planet_list}\n\n"
    
    return chart_path, report

def gradio_interface(birth_date, birth_location, transit_option, custom_transit_date, days_offset):
    """Gradio wrapper function."""
    try:
        if not birth_date:
            return None, "❌ Please enter your birth date."
        
        if not birth_location:
            return None, "❌ Please enter your birth location."
        
        if transit_option == "Today":
            transit_date = datetime.now().strftime('%Y-%m-%d')
        elif transit_option == "Tomorrow":
            transit_date = (datetime.now() + timedelta(days=1)).strftime('%Y-%m-%d')
        elif transit_option == "Custom Date":
            if not custom_transit_date:
                return None, "❌ Please enter a custom transit date."
            transit_date = custom_transit_date
        elif transit_option == "Days from Today":
            transit_date = (datetime.now() + timedelta(days=int(days_offset))).strftime('%Y-%m-%d')
        else:
            transit_date = datetime.now().strftime('%Y-%m-%d')
        
        chart_path, report = generate_rashi_chart(birth_date, birth_location, transit_date)
        return chart_path, report
    except Exception as e:
        return None, f"❌ Error: {str(e)}\n\nPlease check your inputs and try again."

def update_visibility(option):
    """Update visibility of date input fields."""
    return {
        custom_transit_date: gr.update(visible=(option == "Custom Date")),
        days_offset: gr.update(visible=(option == "Days from Today"))
    }

# Create Gradio Interface
with gr.Blocks(title="Vedic Rashi Chart Generator") as demo:
    gr.Markdown("""
    # 🌙 Vedic Rashi Chart Generator
    
    Generate your personalized Vedic astrology Rashi (Moon sign) chart with current planetary transits.
    
    
    """)
    
    with gr.Row():
        with gr.Column():
            birth_date = gr.Textbox(
                label="📅 Birth Date",
                placeholder="YYYY-MM-DD (e.g., 1995-08-15)",
                value=""
            )
            birth_location = gr.Textbox(
                label="📍 Birth Location",
                placeholder="e.g., Mumbai, India",
                value=""
            )
            
            gr.Markdown("### 🕐 Transit Date Options")
            transit_option = gr.Radio(
                choices=["Today", "Tomorrow", "Custom Date", "Days from Today"],
                value="Today",
                label="Select Transit Date Option"
            )
            
            custom_transit_date = gr.Textbox(
                label="Custom Transit Date",
                placeholder="YYYY-MM-DD",
                visible=False
            )
            
            days_offset = gr.Number(
                label="Days from Today",
                value=0,
                visible=False
            )
            
            generate_btn = gr.Button("✨ Generate Rashi Chart", variant="primary")
        
        with gr.Column():
            chart_output = gr.Image(label="📊 Rashi Chart", type="filepath")
            report_output = gr.Markdown(label="📋 Detailed Report")
    
    transit_option.change(
        update_visibility,
        inputs=[transit_option],
        outputs=[custom_transit_date, days_offset]
    )
    
    generate_btn.click(
        gradio_interface,
        inputs=[birth_date, birth_location, transit_option, custom_transit_date, days_offset],
        outputs=[chart_output, report_output]
    )

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