import streamlit as st
import pandas as pd
from datetime import datetime
import pytz
from transformers import pipeline
import math

# Set page config
st.set_page_config(
    page_title="Hindu Astrology Daily Prediction",
    page_icon="🔮",
    layout="wide"
)


# Initialize the text generation pipeline
@st.cache_resource
def load_model():
    return pipeline('text-generation', model='gpt2')


try:
    generator = load_model()
except Exception as e:
    st.error("Error loading the model. Please try again later.")
    st.stop()


def calculate_zodiac_sign(month, day):
    """Calculate zodiac sign based on birth date."""
    if (month == 3 and day >= 21) or (month == 4 and day <= 19):
        return "Aries"
    elif (month == 4 and day >= 20) or (month == 5 and day <= 20):
        return "Taurus"
    elif (month == 5 and day >= 21) or (month == 6 and day <= 20):
        return "Gemini"
    elif (month == 6 and day >= 21) or (month == 7 and day <= 22):
        return "Cancer"
    elif (month == 7 and day >= 23) or (month == 8 and day <= 22):
        return "Leo"
    elif (month == 8 and day >= 23) or (month == 9 and day <= 22):
        return "Virgo"
    elif (month == 9 and day >= 23) or (month == 10 and day <= 22):
        return "Libra"
    elif (month == 10 and day >= 23) or (month == 11 and day <= 21):
        return "Scorpio"
    elif (month == 11 and day >= 22) or (month == 12 and day <= 21):
        return "Sagittarius"
    elif (month == 12 and day >= 22) or (month == 1 and day <= 19):
        return "Capricorn"
    elif (month == 1 and day >= 20) or (month == 2 and day <= 18):
        return "Aquarius"
    else:
        return "Pisces"


def get_element(sign):
    """Get the element associated with a zodiac sign."""
    elements = {
        "Aries": "Fire", "Leo": "Fire", "Sagittarius": "Fire",
        "Taurus": "Earth", "Virgo": "Earth", "Capricorn": "Earth",
        "Gemini": "Air", "Libra": "Air", "Aquarius": "Air",
        "Cancer": "Water", "Scorpio": "Water", "Pisces": "Water"
    }
    return elements.get(sign, "Unknown")


def calculate_nakshatra(birth_date):
    """Calculate Nakshatra based on birth date (simplified calculation)."""
    nakshatras = [
        "Ashwini", "Bharani", "Krittika", "Rohini", "Mrigashira", "Ardra",
        "Punarvasu", "Pushya", "Ashlesha", "Magha", "Purva Phalguni",
        "Uttara Phalguni", "Hasta", "Chitra", "Swati", "Vishakha", "Anuradha",
        "Jyeshtha", "Mula", "Purva Ashadha", "Uttara Ashadha", "Shravana",
        "Dhanishta", "Shatabhisha", "Purva Bhadrapada", "Uttara Bhadrapada", "Revati"
    ]
    # Simplified calculation based on birth date
    day_of_year = birth_date.timetuple().tm_yday
    nakshatra_index = (day_of_year * 27) // 365
    return nakshatras[nakshatra_index]


def generate_prediction(birth_date, birth_time, timezone):
    """Generate prediction based on astrological factors."""
    zodiac_sign = calculate_zodiac_sign(birth_date.month, birth_date.day)
    element = get_element(zodiac_sign)
    nakshatra = calculate_nakshatra(birth_date)

    prompt = f"""
    Astrological reading for a person born under:
    - Zodiac Sign: {zodiac_sign}
    - Element: {element}
    - Nakshatra: {nakshatra}
    - Birth Time: {birth_time.strftime('%H:%M')}

    Based on these positions, today's prediction: 
    """

    try:
        prediction = generator(
            prompt,
            max_length=200,
            num_return_sequences=1,
            temperature=0.7
        )[0]['generated_text']

        return prediction.split("today's prediction: ")[-1].strip()
    except Exception as e:
        return "Unable to generate prediction at this time. Please try again later."


# Streamlit UI
st.title("🔮 Hindu Astrology Daily Prediction")
st.write("Enter your birth details to receive your personalized daily prediction.")

# Input forms
col1, col2 = st.columns(2)

with col1:
    birth_date = st.date_input(
        "Date of Birth",
        min_value=datetime(1900, 1, 1),
        max_value=datetime.now()
    )

with col2:
    birth_time = st.time_input("Time of Birth")

# Time zone selection
timezone = st.selectbox(
    "Select Your Time Zone",
    options=pytz.all_timezones,
    index=pytz.all_timezones.index('UTC')
)

if st.button("Get Prediction"):
    with st.spinner("Calculating astrological factors and generating prediction..."):
        try:
            # Display birth chart information
            st.subheader("Your Astrological Profile")
            zodiac_sign = calculate_zodiac_sign(birth_date.month, birth_date.day)
            element = get_element(zodiac_sign)
            nakshatra = calculate_nakshatra(birth_date)

            profile_data = pd.DataFrame([{
                "Zodiac Sign": zodiac_sign,
                "Element": element,
                "Nakshatra": nakshatra,
                "Birth Time": birth_time.strftime("%H:%M")
            }])

            st.table(profile_data)

            # Generate and display prediction
            st.subheader("Your Daily Prediction")
            prediction = generate_prediction(birth_date, birth_time, timezone)
            st.write(prediction)

            # Add disclaimer
            st.markdown("---")
            st.caption(
                "Disclaimer: This prediction is generated using AI and should be taken "
                "as entertainment only. For accurate astrological readings, please "
                "consult a professional astrologer."
            )

        except Exception as e:
            st.error(f"An error occurred: {str(e)}")