import ollama
from lightweight_charts import Chart

class PatternGenerator:
    def __init__(self):
        self.client = ollama.Client()
        self.code_model = "codellama:latest"
        self.patterns = [
            'head_and_shoulders', 'double_top', 'double_bottom',
            'triangle', 'channel', 'fibonacci', 'wedge', 'flag',
            'pennant', 'cup_and_handle', 'rounding_bottom'
            # Will expand to 300+ patterns
        ]
    
    def generate_all_patterns(self, chart_image, historical_data):
        pattern_charts = []
        for pattern in self.patterns:
            # Generate pattern drawing code using CodeLlama
            code_prompt = f"Generate Lightweight Charts code to draw {pattern} pattern with proper styling and technical analysis markers"
            code_response = self.client.generate(
                model=self.code_model,
                prompt=code_prompt
            )
            
            # Create chart with pattern overlay
            chart = Chart()
            chart.set(code_response)
            chart.add_data(historical_data)
            
            pattern_charts.append(chart)
        
        return pattern_charts