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

# 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")

# 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."

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

# 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 600 words."
    ),
    expected_output="A polished, stakeholder-ready patent insights document with actionable recommendations.",
    agent=writer
)

analyse = Task(
    description=(
        "1. Perform statistical analysis of patent filing trends, innovation hot spots, and growth projections.\n"
        "2. Identify top regions, key players, and technology combinations.\n"
        "3. Generate visualizations such as heatmaps and multi-line charts for trends.\n"
        "4. Provide structured output with fields 'Category' and 'Values' for visualization.\n"
        "5. Collaborate with the Planner and Writer to align on data needs."
    ),
    expected_output="A detailed statistical analysis with actionable insights, heatmaps, and trends.",
    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:
                pdf.add_page()
                pdf.image(chart_path, x=10, y=20, w=180)

        if table_data:
            pdf.add_page()
            pdf.set_font("Arial", size=10)
            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:
            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)

        except Exception as 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 Patent Insights"):
    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.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:
            st.error(f"An error occurred during execution: {e}")