app.py

import streamlit as st
import base64
from datetime import datetime
import json
import os
from anthropic import Anthropic
from PIL import Image
import io
import yaml

def load_anthropic_key():
    """Load Anthropic API key from environment variables"""
    try:
        api_key = os.getenv('ANTHROPIC_API_KEY')
        if not api_key:
            st.error("Anthropic API key not found. Please set ANTHROPIC_API_KEY in your environment variables.")
            return None
        return api_key
    except Exception as e:
        st.error(f"Error loading Anthropic API key: {str(e)}")
        return None

def initialize_anthropic_client():
    """Initialize Anthropic client with API key"""
    api_key = load_anthropic_key()
    if api_key:
        return Anthropic(api_key=api_key)
    return None

def create_prompt_template(patterns, indicators):
    """Creates a structured prompt for the LLM based on the chart and analysis needs"""
    prompt = """You are an expert financial analyst. Please analyze this financial chart (chart type will be detected automatically) and provide insights in the following structured format:

    1. VISUAL ANALYSIS
    - First identify the type of chart (candlestick, line, OHLC, area, etc.)
    - Identify and describe the main trend
    - Note key price levels visible in the chart
    - Describe any significant patterns: {patterns}
    - Comment on volume trends if visible
    - Analyze these technical indicators: {indicators}

    2. TECHNICAL INTERPRETATION
    - Current market structure and trend strength
    - Key support and resistance levels with price points
    - Any visible divergences or convergences
    - Pattern reliability assessment

    3. RISK ANALYSIS
    - Potential risk levels
    - Risk/reward scenarios
    - Warning signs or red flags
    - Market context considerations

    4. ACTIONABLE INSIGHTS
    - Potential trading scenarios
    - Key price targets
    - Suggested stop-loss levels
    - Timeframe considerations

    5. SIMPLIFIED EXPLANATION
    Provide a 2-3 sentence summary in simple terms for novice traders.

    IMPORTANT: Clearly mark this as AI-generated analysis for informational purposes only.
    """
    return prompt.format(
        patterns=', '.join(patterns) if patterns else 'all visible patterns',
        indicators=', '.join(indicators) if indicators else 'visible indicators'
    )

def detect_chart_type(client, image_data):
    """Detect chart type using Claude Vision"""
    try:
        encoded_image = base64.b64encode(image_data).decode('utf-8')
        
        message = client.messages.create(
            model="claude-3-opus-20240229",
            max_tokens=50,
            messages=[{
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "What type of financial chart is this? Choose from: Candlestick, Line, OHLC, Area, or Other. Just respond with one word."
                    },
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": "image/jpeg",
                            "data": encoded_image
                        }
                    }
                ]
            }]
        )
        
        chart_type = message.content[0].text.strip()
        return chart_type
    except Exception as e:
        st.error(f"Error in chart type detection: {str(e)}")
        return "Other"

def analyze_chart_with_claude(client, image_data, prompt, chart_type=None):
    """Analyze chart using Claude Vision"""
    try:
        encoded_image = base64.b64encode(image_data).decode('utf-8')
        
        # If chart type wasn't provided, detect it first
        if not chart_type:
            chart_type = detect_chart_type(client, image_data)
            st.info(f"Detected chart type: {chart_type}")
        
        message = client.messages.create(
            model="claude-3-opus-20240229",
            max_tokens=1000,
            messages=[{
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt.format(chart_type=chart_type)
                    },
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": "image/jpeg",
                            "data": encoded_image
                        }
                    }
                ]
            }]
        )
        
        return message.content[0].text, chart_type
    except Exception as e:
        st.error(f"Error in Claude analysis: {str(e)}")
        return None, None

def continue_analysis_with_claude(client, question, previous_analysis, image_data=None):
    """Continue the analysis based on a follow-up question"""
    try:
        content = [
            {
                "type": "text",
                "text": f"""Previous analysis: {previous_analysis}

User's follow-up question: {question}

Please provide a detailed answer to the follow-up question, maintaining the context of the previous analysis."""
            }
        ]
        
        # Add image to the content if available
        if image_data:
            encoded_image = base64.b64encode(image_data).decode('utf-8')
            content.append({
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/jpeg",
                    "data": encoded_image
                }
            })
        
        message = client.messages.create(
            model="claude-3-opus-20240229",
            max_tokens=1000,
            messages=[{
                "role": "user",
                "content": content
            }]
        )
        
        return message.content[0].text
    except Exception as e:
        st.error(f"Error in follow-up analysis: {str(e)}")
        return None

def get_trading_education(client, concept):
    """Get educational content about trading concepts"""
    try:
        message = client.messages.create(
            model="claude-3-opus-20240229",
            max_tokens=1000,
            messages=[{
                "role": "user",
                "content": f"""Please explain the trading concept '{concept}' in a clear, educational way. Structure your response as follows:

1. Basic Definition
2. How it Works
3. Key Characteristics
4. When to Look for It
5. Trading Implications
6. Common Mistakes to Avoid
7. Real-World Example

If relevant, describe what a typical chart pattern for this concept looks like.
Include any important formulas or calculations if applicable.

Please make this explanation suitable for beginners while also including enough depth for intermediate traders."""
            }]
        )
        
        return message.content[0].text
    except Exception as e:
        st.error(f"Error in getting educational content: {str(e)}")
        return None

def extract_stock_info(analysis_text):
    """Extract stock name and other metadata from analysis text"""
    # This is a simple implementation - can be made more sophisticated
    stock_name = "Unknown"
    try:
        # Look for common stock name patterns
        if "analyzing" in analysis_text.lower():
            words = analysis_text.split()
            for i, word in enumerate(words):
                if word.lower() == "analyzing":
                    stock_name = words[i + 1].strip("(),.:")
    except:
        pass
    return stock_name

def save_chat_history(chat_history, image_data=None, filename=None):
    """Saves chat history and associated image to JSON and image files"""
    if not os.path.exists("chat_histories"):
        os.makedirs("chat_histories")
    if not os.path.exists("chat_images"):
        os.makedirs("chat_images")
        
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # Get the stock name from the latest analysis
    stock_name = "Unknown"
    if chat_history:
        latest_analysis = chat_history[-1]['analysis']
        stock_name = extract_stock_info(latest_analysis)
    
    # Create filename with metadata
    if filename:
        base_filename = filename
    else:
        base_filename = f"{stock_name}_{timestamp}"
    
    # Save image if provided
    image_filename = None
    if image_data:
        image_filename = f"{base_filename}.jpg"
        image_path = os.path.join("chat_images", image_filename)
        with open(image_path, "wb") as f:
            f.write(image_data)
    
    # Add metadata to chat history
    chat_data = {
        'metadata': {
            'stock_name': stock_name,
            'date_created': timestamp,
            'image_file': image_filename
        },
        'conversations': chat_history
    }
    
    # Save chat history
    json_filename = f"{base_filename}.json"
    filepath = os.path.join("chat_histories", json_filename)
    with open(filepath, "w") as f:
        json.dump(chat_data, f)
    
    return json_filename

def load_chat_history(filename):
    """Loads chat history and associated image"""
    filepath = os.path.join("chat_histories", filename)
    with open(filepath, "r") as f:
        chat_data = json.load(f)
        
    # Load associated image if it exists
    image_data = None
    if chat_data.get('metadata', {}).get('image_file'):
        image_path = os.path.join("chat_images", chat_data['metadata']['image_file'])
        if os.path.exists(image_path):
            with open(image_path, "rb") as f:
                image_data = f.read()
                
    return chat_data, image_data

def main():
    st.set_page_config(
        page_title="Stock Chart Assistant",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Initialize Anthropic client
    client = initialize_anthropic_client()
    if not client:
        st.error("Failed to initialize Anthropic client. Please check your API key configuration.")
        return
    
    # Initialize session state
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    if 'current_image' not in st.session_state:
        st.session_state.current_image = None
    if 'current_analysis' not in st.session_state:
        st.session_state.current_analysis = None
    
    # Tab selection
    tab1, tab2 = st.tabs(["Chart Analysis", "Learn Trading"])
    
    with tab1:
        # Initialize variables
        uploaded_file = None
        screenshot_taken = False

        # Sidebar
        with st.sidebar:
            st.title("🚀 Chart Analysis AI")
            upload_option = st.radio(
                "Choose input method:",
                ("Upload Image", "Take Screenshot", "Ask Question"),
                key="analysis_upload_option"  # Added unique key
            )
            
            # File uploader
            if upload_option == "Upload Image":
                uploaded_file = st.file_uploader("Upload your chart", type=["png", "jpg", "jpeg"], key="analysis_file_uploader")
                if uploaded_file:
                    st.session_state.current_image = uploaded_file.getvalue()
            elif upload_option == "Take Screenshot":
                if st.button("Take Screenshot", key="analysis_screenshot_button"):
                    st.info("Feature coming soon! For now, please use the Upload Image option.")
                    screenshot_taken = False
            
            # Analysis Options
            st.subheader("Analysis Options")
            patterns = st.multiselect(
                "Patterns to Look For",
                ["Double Top/Bottom", "Head and Shoulders", "Triangle", 
                 "Flag", "Wedge", "Channel", "Support/Resistance"],
                key="analysis_patterns"
            )
            
            indicators = st.multiselect(
                "Technical Indicators",
                ["Moving Averages", "RSI", "MACD", "Bollinger Bands", 
                 "Volume", "Stochastic", "ADX"],
                key="analysis_indicators"
            )

        # Main content area
        st.title("📈 Stock Chart Analysis Assistant")
        
        # Create two columns for layout
        col1, col2 = st.columns([2, 1])
        
        with col1:
            if upload_option == "Ask Question":
                user_question = st.text_input("What would you like to know about your chart?")
            
            # Display uploaded image
            if uploaded_file is not None:
                st.image(uploaded_file, caption="Uploaded Chart", use_container_width=True)
            
            # Continue chat section
            if st.session_state.current_analysis:
                st.subheader("Continue Analysis")
                follow_up = st.text_input("Ask a follow-up question about this chart:")
                if st.button("Send Follow-up"):
                    if follow_up:
                        with st.spinner("Analyzing..."):
                            follow_up_response = continue_analysis_with_claude(
                                client,
                                follow_up,
                                st.session_state.current_analysis,
                                st.session_state.current_image
                            )
                            if follow_up_response:
                                st.write(follow_up_response)
                                # Add to chat history
                                st.session_state.chat_history.append({
                                    'timestamp': datetime.now().isoformat(),
                                    'question': follow_up,
                                    'analysis': follow_up_response
                                })
            
            if st.button("Analyze"):
                if upload_option == "Ask Question" and user_question:
                    st.info("Question-based analysis feature coming soon!")
                elif uploaded_file is None and not screenshot_taken:
                    st.warning("Please upload an image or take a screenshot first.")
                else:
                    with st.spinner("Analyzing chart..."):
                        # Generate prompt
                        prompt = create_prompt_template(patterns, indicators)
                        
                        if uploaded_file:
                            # Process image and get analysis
                            analysis_result, chart_type = analyze_chart_with_claude(
                                client,
                                uploaded_file.getvalue(),
                                prompt
                            )
                            
                            if analysis_result:
                                # Store current analysis
                                st.session_state.current_analysis = analysis_result
                                
                                # Add to chat history
                                st.session_state.chat_history.append({
                                    'timestamp': datetime.now().isoformat(),
                                    'chart_type': chart_type,
                                    'analysis': analysis_result
                                })
                                
                                # Display analysis
                                st.subheader("Analysis Results")
                                st.write(analysis_result)
                                
                                # Risk warning
                                st.warning(
                                    "⚠️ This analysis is AI-generated and for informational purposes only. "
                                    "Do not make trading decisions solely based on this information."
                                )
        
        with col2:
            st.subheader("Chat History")
            
            # Display chat history
            for chat in st.session_state.chat_history:
                timestamp = datetime.fromisoformat(chat['timestamp']).strftime("%Y-%m-%d %H:%M")
                with st.expander(f"Analysis from {timestamp}"):
                    st.write(chat['analysis'])
                    if 'question' in chat:
                        st.write(f"Follow-up: {chat['question']}")
            
            # Save chat options
            save_name = st.text_input("Save chat as (optional):", key="save_chat_name")
            if st.button("Save Chat", key="save_chat_button"):
                if st.session_state.chat_history:
                    filename = save_chat_history(
                        st.session_state.chat_history,
                        st.session_state.current_image,
                        f"{save_name}.json" if save_name else None
                    )
                    st.success(f"Chat saved as {filename}")
                else:
                    st.warning("No chat history to save.")
    
    with tab2:
        st.title("📚 Learn Trading")
        
        # Search or select trading concept
        concept = st.text_input("Enter a trading concept you'd like to learn about (e.g., 'evening star pattern', 'RSI', 'MACD'):", key="learn_concept")
        if st.button("Learn", key="learn_button"):
            if concept:
                with st.spinner("Getting educational content..."):
                    education_content = get_trading_education(client, concept)
                    if education_content:
                        st.markdown(education_content)

if __name__ == "__main__":
    main()
    # Initialize variables
    uploaded_file = None
    screenshot_taken = False

    # Sidebar
    with st.sidebar:
        st.title("🚀 Chart Analysis AI")
        upload_option = st.radio(
            "Choose input method:",
            ("Upload Image", "Take Screenshot", "Ask Question")
        )
        
        # File uploader
        if upload_option == "Upload Image":
            uploaded_file = st.file_uploader("Upload your chart", type=["png", "jpg", "jpeg"])
        elif upload_option == "Take Screenshot":
            if st.button("Take Screenshot", key="screenshot"):
                st.info("Feature coming soon! For now, please use the Upload Image option.")
                screenshot_taken = False
        
        # Analysis Options
        st.subheader("Analysis Options")
        patterns = st.multiselect(
            "Patterns to Look For",
            ["Double Top/Bottom", "Head and Shoulders", "Triangle", 
             "Flag", "Wedge", "Channel", "Support/Resistance"]
        )
        
        indicators = st.multiselect(
            "Technical Indicators",
            ["Moving Averages", "RSI", "MACD", "Bollinger Bands", 
             "Volume", "Stochastic", "ADX"]
        )

    # Main content area
    st.title("📈 Stock Chart Analysis Assistant")
    
    # Create two columns for layout
    col1, col2 = st.columns([2, 1])
    
    with col1:
        if upload_option == "Ask Question":
            user_question = st.text_input("What would you like to know about your chart?")
        
        # Display uploaded image
        if uploaded_file is not None:
            st.image(uploaded_file, caption="Uploaded Chart", use_container_width=True)
        
        if st.button("Analyze"):
            if upload_option == "Ask Question" and user_question:
                st.info("Question-based analysis feature coming soon!")
            elif uploaded_file is None and not screenshot_taken:
                st.warning("Please upload an image or take a screenshot first.")
            else:
                with st.spinner("Analyzing chart..."):
                    # Generate prompt
                    prompt = create_prompt_template(patterns, indicators)
                    
                    if uploaded_file:
                        # Process image and get analysis
                        analysis_result, chart_type = analyze_chart_with_claude(
                            client,
                            uploaded_file.getvalue(),
                            prompt
                        )
                        
                        if analysis_result:
                            # Add to chat history
                            st.session_state.chat_history.append({
                                'timestamp': datetime.now().isoformat(),
                                'chart_type': chart_type,
                                'analysis': analysis_result
                            })
                            
                            # Display analysis
                            st.subheader("Analysis Results")
                            st.write(analysis_result)
                            
                            # Risk warning
                            st.warning(
                                "⚠️ This analysis is AI-generated and for informational purposes only. "
                                "Do not make trading decisions solely based on this information."
                            )
    
    with col2:
        st.subheader("Chat History")
        
        # Display chat history
        for chat in st.session_state.chat_history:
            with st.expander(f"Analysis from {chat['timestamp'][:16]}"):
                st.write(chat['analysis'])
        
        # Save chat options
        save_name = st.text_input("Save chat as (optional):")
        if st.button("Save Chat"):
            if st.session_state.chat_history:
                filename = save_chat_history(
                    st.session_state.chat_history,
                    f"{save_name}.json" if save_name else None
                )
                st.success(f"Chat saved as {filename}")
            else:
                st.warning("No chat history to save.")

if __name__ == "__main__":
    main()