# to-do: add inventor + assignee
import streamlit as st
from crewai import Agent, Task, Crew
import os
from langchain_groq import ChatGroq
from langchain_openai import ChatOpenAI
from fpdf import FPDF
import pandas as pd
import plotly.express as px
import tempfile
import time
import logging

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Title and Application Introduction
st.title("Patent Strategy and Innovation Consultant")
st.sidebar.write(
    "This application uses AI to provide actionable insights and comprehensive analysis for patent-related strategies."
)

# User Input Section
st.sidebar.header("User Inputs")
patent_area = st.text_input("Enter Patent Technology Area", value="Transparent Antennas for Windshields")
stakeholder = st.text_input("Enter Stakeholder", value="Patent Attorneys")

# Initialize LLM
llm = None

# Model Selection
#st.header("Model Selection")
model_choice = st.radio("Select LLM", ["OpenAI Model", "Groq-based LLM"], index=0, horizontal=True)


# API Key Validation and LLM Initialization
groq_api_key = os.getenv("GROQ_API_KEY")
openai_api_key = os.getenv("OPENAI_API_KEY")

#llm = ChatGroq(groq_api_key=os.getenv("GROQ_API_KEY"), model="groq/llama-3.3-70b-versatile")

if model_choice == "Groq-based LLM":
    if not groq_api_key:
        st.error("Groq API key is missing. Please set the GROQ_API_KEY environment variable.")
        llm = None
    else:
        llm = ChatGroq(groq_api_key=groq_api_key, model="groq/llama-3.3-70b-versatile")
elif model_choice == "OpenAI Model":
    if not openai_api_key:
        st.error("OpenAI API key is missing. Please set the OPENAI_API_KEY environment variable.")
        llm = None
    else:
        llm = ChatOpenAI(api_key=openai_api_key, model="gpt-4o")

# Advanced Options
st.sidebar.header("Advanced Options")
enable_advanced_analysis = st.sidebar.checkbox("Enable Advanced Analysis", value=True)
enable_custom_visualization = st.sidebar.checkbox("Enable Custom Visualizations", value=True)

# Agent Customization
st.sidebar.header("Agent Customization")
with st.sidebar.expander("Customize Agent Goals", expanded=False):
    enable_customization = st.checkbox("Enable Custom Goals")
    if enable_customization:
        planner_goal = st.text_area(
            "Planner Goal",
            value="Research trends in patent filings and technological innovation, identify key players, and provide strategic recommendations."
        )
        writer_goal = st.text_area(
            "Writer Goal",
            value="Craft a professional insights document summarizing trends, strategies, and actionable outcomes for stakeholders."
        )
        analyst_goal = st.text_area(
            "Analyst Goal",
            value="Perform detailed statistical analysis of patent filings, growth trends, and innovation distribution."
        )
    else:
        planner_goal = "Research trends in patent filings and technological innovation, identify key players, and provide strategic recommendations."
        writer_goal = "Craft a professional insights document summarizing trends, strategies, and actionable outcomes for stakeholders."
        analyst_goal = "Perform detailed statistical analysis of patent filings, growth trends, and innovation distribution."


# Agent Definitions
planner = Agent(
    role="Patent Research Consultant",
    goal=planner_goal,
    backstory=(
        "You're tasked with researching {topic} patents and identifying key trends and players. Your work supports the Patent Writer and Data Analyst."
    ),
    allow_delegation=False,
    verbose=True,
    llm=llm
)

writer = Agent(
    role="Patent Insights Writer",
    goal=writer_goal,
    backstory=(
        "Using the research from the Planner and data from the Analyst, craft a professional document summarizing patent insights for {stakeholder}."
    ),
    allow_delegation=False,
    verbose=True,
    llm=llm
)

analyst = Agent(
    role="Patent Data Analyst",
    goal=analyst_goal,
    backstory=(
        "Analyze patent filing data and innovation trends in {topic} to provide statistical insights. Your analysis will guide the Writer's final report."
    ),
    allow_delegation=False,
    verbose=True,
    llm=llm
)

# Task Definitions
plan = Task(
    description=(
        "1. Research recent trends in {topic} patent filings and innovation.\n"
        "2. Identify key players and emerging technologies.\n"
        "3. Provide recommendations for stakeholders on strategic directions.\n"
        "4. Identify key statistics such as top regions, top players, and hot areas of innovation.\n"
        "5. Limit the output to 500 words."
    ),
    expected_output="A research document with structured insights, strategic recommendations, and key statistics.",
    agent=planner
)

write = Task(
    description=(
        "1. Use the Planner's and Analyst's outputs to craft a professional patent insights document.\n"
        "2. Include key findings, visual aids, and actionable strategies.\n"
        "3. Suggest strategic directions and highlight untapped innovation areas.\n"
        "4. Incorporate summarized tables for key statistics and example inventions.\n"
        "5. Limit the document to 650 words."
    ),
    expected_output="A polished, stakeholder-ready patent insights document with actionable recommendations.",
    agent=writer
)

#analyse = Task(
#    description=(
#        "1. Perform a detailed statistical analysis of patent filing trends, innovation hot spots, and growth projections.\n"
#        "2. Identify and list the top regions, top assignees, prominent inventors, and leading companies in the transparent antenna industry.\n"
#        "3. Highlight regional innovation trends and technology distribution across different geographies.\n"
#        "4. Generate structured output in a list of dictionaries with 'Category' and 'Values' fields for each analysis area. "
#        "This structured output will be utilized to create visualizations later on.\n"
#        "5. Example Output Format:\n"
#        "[{'Category': 'Top Regions', 'Values': ['North America', 'Europe', 'Asia-Pacific']},\n"
#        " {'Category': 'Top Assignees', 'Values': ['Company A', 'Company B']},\n"
#        " {'Category': 'Top Inventors', 'Values': ['Inventor X', 'Inventor Y']},\n"
#        " {'Category': 'Hot Areas of Innovation', 'Values': ['5G Integration', 'Graphene Antennas']}]\n"
#        "6. Collaborate with the Planner and Writer to align on data needs."
#    ),
#    expected_output="A structured list of dictionaries covering regions, assignees, inventors, and innovation trends for visualization and tabular reporting.",
#    agent=analyst
#)

analyse = Task(
    description=(
        "1. Conduct a comprehensive statistical analysis of patent filing trends, innovation hot spots, and future growth projections in the transparent antenna industry.\n"
        "2. Identify and rank the top regions, leading assignees, prominent inventors, and key companies driving innovation.\n"
        "3. Highlight regional innovation trends and the distribution of emerging technologies across different geographies.\n"
        "4. Provide actionable insights and strategic recommendations based on the data.\n"
        "5. Generate structured data in a list of dictionaries with 'Category' and 'Values' fields for each analysis area.\n"
        "   - 'Values' must be structured to support visualizations and tables:\n"
        "     a) Use a dictionary for counts (e.g., {'Region A': 120, 'Region B': 95}),\n"
        "     b) Use a list for qualitative data (e.g., ['Technology A', 'Technology B']),\n"
        "     c) Use descriptive text for strategic insights.\n"
        "6. Example Output Format:\n"
        "[\n"
        "  {'Category': 'Top Regions', 'Values': {'North America': 120, 'Europe': 95, 'Asia-Pacific': 85}},\n"
        "  {'Category': 'Top Assignees', 'Values': {'Company A': 40, 'Company B': 35}},\n"
        "  {'Category': 'Emerging Technologies', 'Values': ['Graphene Antennas', '5G Integration']},\n"
        "  {'Category': 'Strategic Insights', 'Values': 'Collaborations between automotive and material science industries are accelerating innovation.'}\n"
        "]\n"
        "7. Do NOT generate any visualizations directly. Focus only on delivering structured data that can be used to create visualizations and tables.\n"
        "8. Collaborate with the Planner and Writer to align analysis with reporting needs."
    ),
    expected_output="A structured dataset with numeric, list-based, and descriptive data, ready for visualizations and tabular presentation.",
    agent=analyst
)


crew = Crew(
    agents=[planner, analyst, writer],
    tasks=[plan, analyse, write],
    verbose=True
)

# PDF Report Generation
def generate_pdf_report(result, charts=None, table_data=None, metadata=None):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
        pdf = FPDF()
        pdf.add_page()
        pdf.set_font("Arial", size=12)
        pdf.set_auto_page_break(auto=True, margin=15)

        pdf.set_font("Arial", size=16, style="B")
        pdf.cell(200, 10, txt="Patent Strategy and Innovation Report", ln=True, align="C")
        pdf.ln(10)

        if metadata:
            pdf.set_font("Arial", size=10)
            for key, value in metadata.items():
                pdf.cell(200, 10, txt=f"{key}: {value}", ln=True)

        pdf.set_font("Arial", size=12)
        pdf.multi_cell(0, 10, txt=result)

        if charts:
            for chart_path in charts:
                try:
                    pdf.add_page()
                    pdf.image(chart_path, x=10, y=20, w=180)
                    logging.info(f"Successfully included chart: {chart_path}")
                except Exception as e:
                    logging.error(f"Failed to include chart in PDF: {chart_path}. Error: {e}")

        if table_data:
            pdf.add_page()
            pdf.set_font("Arial", size=10)
            pdf.cell(200, 10, txt="Consolidated Table:", ln=True, align="L")
            for row in table_data:
                pdf.cell(200, 10, txt=str(row), ln=True)

        pdf.output(temp_pdf.name)
        return temp_pdf.name

# Data Validation
def validate_analyst_output(analyst_output):
    if not analyst_output:
        st.warning("No data available for analysis.")
        return None
    if not isinstance(analyst_output, list) or not all(isinstance(item, dict) for item in analyst_output):
        st.warning("Analyst output must be a list of dictionaries.")
        return None
    required_keys = {'Category', 'Values'}
    if not all(required_keys.issubset(item.keys()) for item in analyst_output):
        st.warning(f"Each dictionary must contain keys: {required_keys}")
        return None
    return analyst_output

# Visualization and Table Display
def create_visualizations(analyst_output):
    chart_paths = []
    validated_data = validate_analyst_output(analyst_output)
    if validated_data:
        data = pd.DataFrame(validated_data)
        try:
            if data.empty:
                raise ValueError("Data for visualizations is empty.")

            bar_chart = px.bar(data, x="Category", y="Values", title="Patent Trends by Category")
            st.plotly_chart(bar_chart)
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_chart:
                bar_chart.write_image(temp_chart.name)
                chart_paths.append(temp_chart.name)

            pie_chart = px.pie(data, names="Category", values="Values", title="Category Distribution")
            st.plotly_chart(pie_chart)
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_chart:
                pie_chart.write_image(temp_chart.name)
                chart_paths.append(temp_chart.name)

            heatmap_chart = px.density_heatmap(data, x="Category", y="Values", title="Regional Patent Density")
            st.plotly_chart(heatmap_chart)
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_chart:
                heatmap_chart.write_image(temp_chart.name)
                chart_paths.append(temp_chart.name)

            multi_line_chart = px.line(data, x="Category", y="Values", title="Trends Over Time")
            st.plotly_chart(multi_line_chart)
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_chart:
                multi_line_chart.write_image(temp_chart.name)
                chart_paths.append(temp_chart.name)

        except Exception as e:
            logging.error(f"Error generating visualization: {e}")
            st.error(f"Error generating visualization: {e}")
    return chart_paths

def display_table(analyst_output):
    table_data = []
    validated_data = validate_analyst_output(analyst_output)
    if validated_data:
        data = pd.DataFrame(validated_data)
        st.dataframe(data)
        table_data = data.to_dict(orient="records")
    return table_data

# Main Execution Block
if st.button("Generate Insights"):
    if llm is None:
        st.error("Cannot proceed without a valid API key for the selected model.")
    else:
        with st.spinner('Processing...'):
            try:
                start_time = time.time()
                results = crew.kickoff(inputs={"topic": patent_area, "stakeholder": stakeholder})
                elapsed_time = time.time() - start_time

                writer_output = getattr(results.tasks_output[2], "raw", "No details available.")
                if writer_output:
                    st.markdown("### Final Report")
                    st.write(writer_output)
                else:
                    st.warning("No final report available.")

                with st.expander("Explore Detailed Insights"):
                    tab1, tab2 = st.tabs(["Planner's Insights", "Analyst's Analysis"])

                    with tab1:
                        planner_output = getattr(results.tasks_output[0], "raw", "No details available.")
                        st.write(planner_output)

                    with tab2:
                        analyst_output = getattr(results.tasks_output[1], "raw", "No details available.")
                        st.subheader("Raw Analyst Output")
                        st.write(analyst_output)

                        charts = []
                        if enable_advanced_analysis:
                            charts = create_visualizations(analyst_output)

                        table_data = display_table(analyst_output)

                st.success(f"Analysis completed in {elapsed_time:.2f} seconds.")
                pdf_path = generate_pdf_report(
                    writer_output,
                    charts=charts,
                    table_data=table_data,
                    metadata={"Technology Area": patent_area, "Stakeholder": stakeholder}
                )
                with open(pdf_path, "rb") as report_file:
                    st.download_button("Download Report", data=report_file, file_name="Patent_Strategy_Report.pdf")

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