#!/usr/bin/env python3
"""
Treatment Timeline Generator
Generates visual treatment timelines from treatment plan files.
"""

import sys
import re
import argparse
from pathlib import Path
from datetime import datetime, timedelta
from typing import List, Dict, Tuple

# Try to import matplotlib, but make it optional
try:
    import matplotlib.pyplot as plt
    import matplotlib.dates as mdates
    from matplotlib.patches import Rectangle
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False


def extract_timeline_info(content: str) -> Dict[str, List[Tuple[str, str]]]:
    """
    Extract timeline and schedule information from treatment plan.
    Returns dict with phases, appointments, milestones.
    """
    timeline_data = {
        'phases': [],
        'appointments': [],
        'milestones': []
    }
    
    # Extract treatment phases
    # Look for patterns like "Week 1-4: Description" or "Months 1-3: Description"
    phase_patterns = [
        r'(Week[s]?\s*\d+[-–]\d+|Month[s]?\s*\d+[-–]\d+)[:\s]+([^\n]+)',
        r'(POD\s*\d+[-–]\d+)[:\s]+([^\n]+)',
        r'(\d+[-–]\d+\s*week[s]?)[:\s]+([^\n]+)'
    ]
    
    for pattern in phase_patterns:
        matches = re.findall(pattern, content, re.IGNORECASE)
        for timeframe, description in matches:
            timeline_data['phases'].append((timeframe.strip(), description.strip()))
    
    # Extract appointments
    # Look for patterns like "Week 2: Visit" or "Month 3: Follow-up"
    apt_patterns = [
        r'(Week\s*\d+|Month\s*\d+|POD\s*\d+)[:\s]+(Visit|Appointment|Follow-up|Check-up|Consultation)([^\n]*)',
        r'(Every\s+\d+\s+\w+)[:\s]+(Visit|Appointment|therapy|session)([^\n]*)'
    ]
    
    for pattern in apt_patterns:
        matches = re.findall(pattern, content, re.IGNORECASE)
        for timeframe, visit_type, details in matches:
            timeline_data['appointments'].append((timeframe.strip(), f"{visit_type}{details}".strip()))
    
    # Extract milestones/assessments
    # Look for "reassessment", "goal evaluation", "milestone" mentions
    milestone_patterns = [
        r'(Week\s*\d+|Month\s*\d+)[:\s]+(reassess|evaluation|assessment|milestone)([^\n]*)',
        r'(\w+\s*\d+)[:\s]+(HbA1c|labs?|imaging|test)([^\n]*)'
    ]
    
    for pattern in milestone_patterns:
        matches = re.findall(pattern, content, re.IGNORECASE)
        for timeframe, event_type, details in matches:
            timeline_data['milestones'].append((timeframe.strip(), f"{event_type}{details}".strip()))
    
    return timeline_data


def parse_timeframe_to_days(timeframe: str) -> Tuple[int, int]:
    """
    Parse timeframe string to start and end days.
    Examples: "Week 1-4" -> (0, 28), "Month 3" -> (60, 90)
    """
    timeframe = timeframe.lower()
    
    # Week patterns
    if 'week' in timeframe:
        weeks = re.findall(r'\d+', timeframe)
        if len(weeks) == 2:
            start_week = int(weeks[0])
            end_week = int(weeks[1])
            return ((start_week - 1) * 7, end_week * 7)
        elif len(weeks) == 1:
            week = int(weeks[0])
            return ((week - 1) * 7, week * 7)
    
    # Month patterns
    if 'month' in timeframe:
        months = re.findall(r'\d+', timeframe)
        if len(months) == 2:
            start_month = int(months[0])
            end_month = int(months[1])
            return ((start_month - 1) * 30, end_month * 30)
        elif len(months) == 1:
            month = int(months[0])
            return ((month - 1) * 30, month * 30)
    
    # POD (post-operative day) patterns
    if 'pod' in timeframe:
        days = re.findall(r'\d+', timeframe)
        if len(days) == 2:
            return (int(days[0]), int(days[1]))
        elif len(days) == 1:
            day = int(days[0])
            return (day, day + 1)
    
    # Default fallback
    return (0, 7)


def create_text_timeline(timeline_data: Dict, output_file: Path = None):
    """Create a text-based timeline representation."""
    
    lines = []
    lines.append("="*70)
    lines.append("TREATMENT TIMELINE")
    lines.append("="*70)
    
    # Treatment phases
    if timeline_data['phases']:
        lines.append("\nTREATMENT PHASES:")
        lines.append("-"*70)
        for timeframe, description in timeline_data['phases']:
            lines.append(f"{timeframe:20s} | {description}")
    
    # Appointments
    if timeline_data['appointments']:
        lines.append("\nSCHEDULED APPOINTMENTS:")
        lines.append("-"*70)
        for timeframe, details in timeline_data['appointments']:
            lines.append(f"{timeframe:20s} | {details}")
    
    # Milestones
    if timeline_data['milestones']:
        lines.append("\nMILESTONES & ASSESSMENTS:")
        lines.append("-"*70)
        for timeframe, event in timeline_data['milestones']:
            lines.append(f"{timeframe:20s} | {event}")
    
    lines.append("\n" + "="*70)
    
    # Output
    output_text = "\n".join(lines)
    
    if output_file:
        with open(output_file, 'w') as f:
            f.write(output_text)
        print(f"\nText timeline saved to: {output_file}")
    else:
        print(output_text)
    
    return output_text


def create_visual_timeline(timeline_data: Dict, output_file: Path, start_date: str = None):
    """Create a visual Gantt-chart style timeline (requires matplotlib)."""
    
    if not HAS_MATPLOTLIB:
        print("Error: matplotlib not installed. Install with: pip install matplotlib", file=sys.stderr)
        print("Generating text timeline instead...", file=sys.stderr)
        text_output = output_file.with_suffix('.txt')
        create_text_timeline(timeline_data, text_output)
        return
    
    # Parse start date
    if start_date:
        try:
            start = datetime.strptime(start_date, '%Y-%m-%d')
        except ValueError:
            print(f"Invalid date format: {start_date}. Using today.", file=sys.stderr)
            start = datetime.now()
    else:
        start = datetime.now()
    
    # Prepare data for plotting
    phases = []
    for timeframe, description in timeline_data['phases']:
        start_day, end_day = parse_timeframe_to_days(timeframe)
        phases.append({
            'name': f"{timeframe}: {description[:40]}",
            'start': start + timedelta(days=start_day),
            'end': start + timedelta(days=end_day),
            'type': 'phase'
        })
    
    # Add appointments as events
    events = []
    for timeframe, details in timeline_data['appointments']:
        start_day, _ = parse_timeframe_to_days(timeframe)
        events.append({
            'name': f"{timeframe}: {details[:40]}",
            'date': start + timedelta(days=start_day),
            'type': 'appointment'
        })
    
    # Add milestones
    for timeframe, event in timeline_data['milestones']:
        start_day, _ = parse_timeframe_to_days(timeframe)
        events.append({
            'name': f"{timeframe}: {event[:40]}",
            'date': start + timedelta(days=start_day),
            'type': 'milestone'
        })
    
    # Create figure
    fig, ax = plt.subplots(figsize=(12, 8))
    
    # Plot phases as horizontal bars
    y_position = len(phases) + len(events)
    
    for i, phase in enumerate(phases):
        duration = (phase['end'] - phase['start']).days
        ax.barh(y_position - i, duration, left=mdates.date2num(phase['start']),
                height=0.6, color='steelblue', alpha=0.7, edgecolor='black')
        ax.text(mdates.date2num(phase['start']) + duration/2, y_position - i,
                phase['name'], va='center', ha='center', fontsize=9, color='white', weight='bold')
    
    # Plot events as markers
    event_y = y_position - len(phases) - 1
    
    for i, event in enumerate(events):
        marker = 'o' if event['type'] == 'appointment' else 's'
        color = 'green' if event['type'] == 'appointment' else 'orange'
        ax.plot(mdates.date2num(event['date']), event_y - i, marker=marker,
                markersize=10, color=color, markeredgecolor='black')
        ax.text(mdates.date2num(event['date']) + 2, event_y - i, event['name'],
                va='center', ha='left', fontsize=8)
    
    # Format x-axis as dates
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))
    ax.xaxis.set_major_locator(mdates.MonthLocator())
    plt.xticks(rotation=45, ha='right')
    
    # Labels and title
    ax.set_xlabel('Date', fontsize=12, weight='bold')
    ax.set_title('Treatment Plan Timeline', fontsize=14, weight='bold', pad=20)
    ax.set_yticks([])
    ax.grid(axis='x', alpha=0.3, linestyle='--')
    
    # Legend
    from matplotlib.lines import Line2D
    legend_elements = [
        Rectangle((0, 0), 1, 1, fc='steelblue', alpha=0.7, edgecolor='black', label='Treatment Phase'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor='green', markersize=10,
               markeredgecolor='black', label='Appointment'),
        Line2D([0], [0], marker='s', color='w', markerfacecolor='orange', markersize=10,
               markeredgecolor='black', label='Milestone/Assessment')
    ]
    ax.legend(handles=legend_elements, loc='upper right', framealpha=0.9)
    
    plt.tight_layout()
    
    # Save
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"\nVisual timeline saved to: {output_file}")
    
    # Close plot
    plt.close()


def main():
    parser = argparse.ArgumentParser(
        description='Generate treatment timeline visualization',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate text timeline
  python timeline_generator.py --plan my_plan.tex

  # Generate visual timeline (requires matplotlib)
  python timeline_generator.py --plan my_plan.tex --output timeline.png --visual

  # Specify start date for visual timeline
  python timeline_generator.py --plan my_plan.tex --output timeline.pdf --visual --start 2025-02-01

Output formats:
  Text: .txt
  Visual: .png, .pdf, .svg (requires matplotlib)

Note: Visual timeline generation requires matplotlib.
  Install with: pip install matplotlib
        """
    )
    
    parser.add_argument(
        '--plan',
        type=Path,
        required=True,
        help='Treatment plan file to analyze (.tex format)'
    )
    
    parser.add_argument(
        '--output',
        type=Path,
        help='Output file (default: timeline.txt or timeline.png if --visual)'
    )
    
    parser.add_argument(
        '--visual',
        action='store_true',
        help='Generate visual timeline (requires matplotlib)'
    )
    
    parser.add_argument(
        '--start',
        help='Start date for timeline (YYYY-MM-DD format, default: today)'
    )
    
    args = parser.parse_args()
    
    # Check plan file exists
    if not args.plan.exists():
        print(f"Error: File not found: {args.plan}", file=sys.stderr)
        sys.exit(1)
    
    # Read plan
    try:
        with open(args.plan, 'r', encoding='utf-8') as f:
            content = f.read()
    except Exception as e:
        print(f"Error reading file: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Extract timeline information
    print("Extracting timeline information from treatment plan...")
    timeline_data = extract_timeline_info(content)
    
    # Check if any timeline info found
    total_items = (len(timeline_data['phases']) +
                   len(timeline_data['appointments']) +
                   len(timeline_data['milestones']))
    
    if total_items == 0:
        print("\nWarning: No timeline information detected in treatment plan.", file=sys.stderr)
        print("The plan may not contain structured timeline/schedule sections.", file=sys.stderr)
        print("\nTip: Include sections with timeframes like:", file=sys.stderr)
        print("  - Week 1-4: Initial phase", file=sys.stderr)
        print("  - Month 3: Follow-up visit", file=sys.stderr)
        sys.exit(1)
    
    print(f"Found {len(timeline_data['phases'])} phase(s), "
          f"{len(timeline_data['appointments'])} appointment(s), "
          f"{len(timeline_data['milestones'])} milestone(s)")
    
    # Determine output file
    if not args.output:
        if args.visual:
            args.output = Path('timeline.png')
        else:
            args.output = Path('timeline.txt')
    
    # Generate timeline
    if args.visual:
        create_visual_timeline(timeline_data, args.output, args.start)
    else:
        create_text_timeline(timeline_data, args.output)
    
    print(f"\nTimeline generation complete!")


if __name__ == '__main__':
    main()