Spaces:

projectsorg
/

digital-twin-backend

Sleeping

Added cleanup
3a134c0 10 months ago
26.4 kB
"""
Flask API server for the policy engine.
Endpoints for generating policies, applying them, and analyzing impacts.
"""
from flask import Flask, request, jsonify
from flask_cors import CORS
from datetime import datetime
import json
import traceback

from policy_engine import PolicyEngine, get_graph_context_from_file
from graph_engine import GraphState, ImpactAnalyzer
from health_analyzer import HealthImpactAnalyzer
from explainability import generate_policy_explanation
from explainability import generate_policy_explanation
from aqi import register_aqi_routes
from emission_forecast import register_emission_routes
import config
import os


# ============================================================================
# SETUP
# ============================================================================

app = Flask(__name__)
CORS(app)  # Enable CORS for frontend requests

# Initialize engines
policy_engine = PolicyEngine()
health_analyzer = HealthImpactAnalyzer()

# Register AQI routes
register_aqi_routes(app)
register_emission_routes(app)
from aqi_history import register_aqi_history_routes
register_aqi_history_routes(app)

# Pre-initialize AQI Data & Model to prevent timeout on first request
from aqi_history import get_aqi_history
print("Initializing AQI History & Model...")
get_aqi_history()
from aqi_map import generate_heatmap_html, generate_hotspots_html, generate_forecast_hotspots_html, render_map_to_png
from emission_map import generate_emission_heatmap_html, generate_emission_hotspots_html, generate_forecast_emission_hotspots_html, render_map_to_png as render_emission_map_to_png
import threading

# Ensure map directory exists
STATIC_DIR = os.path.join(os.getcwd(), 'static')
os.makedirs(STATIC_DIR, exist_ok=True)
HEATMAP_PATH = os.path.join(STATIC_DIR, 'aqi_map_heatmap.png')
HOTSPOTS_PATH = os.path.join(STATIC_DIR, 'aqi_map_hotspots.png')

# Emission map paths
EMISSION_HEATMAP_PATH = os.path.join(STATIC_DIR, 'emission_map_heatmap.png')
EMISSION_HOTSPOTS_PATH = os.path.join(STATIC_DIR, 'emission_map_hotspots.png')

# Default map path (alias to heatmap for backward compat)
MAP_IMAGE_PATH = HEATMAP_PATH 

# ============================================================================
# STARTUP CLEANUP - Remove orphaned temp files from crashed runs
# ============================================================================
def cleanup_temp_files():
    """Remove orphaned temp_*.html files from previous crashed runs."""
    import glob
    backend_dir = os.path.dirname(os.path.abspath(__file__))
    patterns = [
        os.path.join(backend_dir, 'temp_map_*.html'),
        os.path.join(backend_dir, 'temp_emission_map_*.html'),
    ]
    cleaned = 0
    for pattern in patterns:
        for temp_file in glob.glob(pattern):
            try:
                os.remove(temp_file)
                cleaned += 1
                print(f"Cleaned up orphaned temp file: {os.path.basename(temp_file)}")
            except Exception as e:
                print(f"Failed to cleanup {temp_file}: {e}")
    if cleaned > 0:
        print(f"Startup cleanup: Removed {cleaned} orphaned temp file(s)")

# Run cleanup on startup
cleanup_temp_files()

@app.route('/api/aqi-map', methods=['GET'])
def get_aqi_map_html():
    """Returns the Interactive Heatmap HTML (Default)"""
    return generate_heatmap_html()

@app.route('/api/aqi-map/heatmap', methods=['GET'])
def get_heatmap_html():
    """Explicit endpoint for Heatmap HTML"""
    return generate_heatmap_html()

@app.route('/api/aqi-map/hotspots', methods=['GET'])
def get_hotspots_html():
    """Returns the Interactive Hotspots HTML. Accepts optional 'year' param for forecast."""
    year = request.args.get('year', type=int)
    
    # Validate year if provided
    if year and year in [2026, 2027, 2028]:
        return generate_forecast_hotspots_html(year)
    
    return generate_hotspots_html()

@app.route('/api/aqi-map.png', methods=['GET'])
def get_aqi_map_png():
    """Default map image (Heatmap)"""
    return get_heatmap_png()

@app.route('/api/aqi-map/heatmap.png', methods=['GET'])
def get_heatmap_png():
    """Returns the Heatmap Grid PNG"""
    from flask import send_file
    
    # Simple caching/regeneration logic
    if not os.path.exists(HEATMAP_PATH) or request.args.get('refresh'):
        html = generate_heatmap_html()
        render_map_to_png(html, HEATMAP_PATH)
        
    return send_file(HEATMAP_PATH, mimetype='image/png')

@app.route('/api/aqi-map/hotspots.png', methods=['GET'])
def get_hotspots_png():
    """Returns the Hotspots PNG. Accepts optional 'year' param for forecast (2026-2028)."""
    from flask import send_file
    
    year = request.args.get('year', type=int)
    
    # Validate year if provided
    if year and year not in [2026, 2027, 2028]:
        year = None
    
    if year:
        # Generate forecast-based hotspots for the selected year
        cache_path = os.path.join(STATIC_DIR, f'aqi_map_hotspots_{year}.png')
        if not os.path.exists(cache_path) or request.args.get('refresh'):
            html = generate_forecast_hotspots_html(year)
            render_map_to_png(html, cache_path)
        return send_file(cache_path, mimetype='image/png')
    else:
        # Original baseline hotspots
        if not os.path.exists(HOTSPOTS_PATH) or request.args.get('refresh'):
            html = generate_hotspots_html()
            render_map_to_png(html, HOTSPOTS_PATH)
        return send_file(HOTSPOTS_PATH, mimetype='image/png')

# ============================================================================
# EMISSION MAPS - CO2 Visualization
# ============================================================================

@app.route('/api/emission-map', methods=['GET'])
def get_emission_map_html():
    """Returns the Interactive Emission Heatmap HTML (Default)"""
    return generate_emission_heatmap_html()

@app.route('/api/emission-map/heatmap', methods=['GET'])
def get_emission_heatmap_html():
    """Explicit endpoint for Emission Heatmap HTML"""
    return generate_emission_heatmap_html()

@app.route('/api/emission-map/hotspots', methods=['GET'])
def get_emission_hotspots_html():
    """Returns the Interactive Emission Hotspots HTML. Accepts optional 'year' param."""
    year = request.args.get('year', type=int)
    
    if year and year in [2026, 2027, 2028]:
        return generate_forecast_emission_hotspots_html(year)
    
    return generate_emission_hotspots_html()

@app.route('/api/emission-map/heatmap.png', methods=['GET'])
def get_emission_heatmap_png():
    """Returns the Emission Heatmap Grid PNG"""
    from flask import send_file
    
    if not os.path.exists(EMISSION_HEATMAP_PATH) or request.args.get('refresh'):
        html = generate_emission_heatmap_html()
        render_emission_map_to_png(html, EMISSION_HEATMAP_PATH)
        
    return send_file(EMISSION_HEATMAP_PATH, mimetype='image/png')

@app.route('/api/emission-map/hotspots.png', methods=['GET'])
def get_emission_hotspots_png():
    """Returns the Emission Hotspots PNG. Accepts optional 'year' param (2026-2028)."""
    from flask import send_file
    
    year = request.args.get('year', type=int)
    
    if year and year not in [2026, 2027, 2028]:
        year = None
    
    if year:
        cache_path = os.path.join(STATIC_DIR, f'emission_map_hotspots_{year}.png')
        if not os.path.exists(cache_path) or request.args.get('refresh'):
            html = generate_forecast_emission_hotspots_html(year)
            render_emission_map_to_png(html, cache_path)
        return send_file(cache_path, mimetype='image/png')
    else:
        if not os.path.exists(EMISSION_HOTSPOTS_PATH) or request.args.get('refresh'):
            html = generate_emission_hotspots_html()
            render_emission_map_to_png(html, EMISSION_HOTSPOTS_PATH)
        return send_file(EMISSION_HOTSPOTS_PATH, mimetype='image/png')

# ============================================================================
# SECTOR-SPECIFIC EMISSION MAPS
# ============================================================================

VALID_SECTORS = ['Industry', 'Transport', 'Power', 'Residential', 'Aviation', 'Commercial']

@app.route('/api/emission-map/sector/heatmap', methods=['GET'])
def get_sector_heatmap_html():
    """Returns sector-specific emission heatmap HTML."""
    from emission_map import generate_sector_heatmap_html
    
    sector = request.args.get('sector', 'Industry')
    if sector not in VALID_SECTORS:
        sector = 'Industry'
    
    return generate_sector_heatmap_html(sector)

@app.route('/api/emission-map/sector/heatmap.png', methods=['GET'])
def get_sector_heatmap_png():
    """Returns sector-specific emission heatmap PNG."""
    from flask import send_file
    from emission_map import generate_sector_heatmap_html, render_map_to_png
    
    sector = request.args.get('sector', 'Industry')
    if sector not in VALID_SECTORS:
        sector = 'Industry'
    
    cache_path = os.path.join(STATIC_DIR, f'sector_heatmap_{sector.lower()}.png')
    if not os.path.exists(cache_path) or request.args.get('refresh'):
        html = generate_sector_heatmap_html(sector)
        render_map_to_png(html, cache_path)
    
    return send_file(cache_path, mimetype='image/png')

@app.route('/api/emission-map/sector/hotspots', methods=['GET'])
def get_sector_hotspots_html():
    """Returns sector-specific emission hotspots HTML."""
    from emission_map import generate_sector_hotspots_html
    
    sector = request.args.get('sector', 'Industry')
    year = request.args.get('year', type=int)
    
    if sector not in VALID_SECTORS:
        sector = 'Industry'
    if year and year not in [2026, 2027, 2028]:
        year = None
    
    return generate_sector_hotspots_html(sector, year)

@app.route('/api/emission-map/sector/hotspots.png', methods=['GET'])
def get_sector_hotspots_png():
    """Returns sector-specific emission hotspots PNG."""
    from flask import send_file
    from emission_map import generate_sector_hotspots_html, render_map_to_png
    
    sector = request.args.get('sector', 'Industry')
    year = request.args.get('year', type=int)
    
    if sector not in VALID_SECTORS:
        sector = 'Industry'
    if year and year not in [2026, 2027, 2028]:
        year = None
    
    cache_key = f'sector_hotspots_{sector.lower()}'
    if year:
        cache_key += f'_{year}'
    cache_path = os.path.join(STATIC_DIR, f'{cache_key}.png')
    
    if not os.path.exists(cache_path) or request.args.get('refresh'):
        html = generate_sector_hotspots_html(sector, year)
        render_map_to_png(html, cache_path)
    
    return send_file(cache_path, mimetype='image/png')

# Background task to refresh map periodically (optional)
def refresh_map_periodically():
    while True:
        try:
            print("Refreshing AQI Map image...")
            html = generate_aqi_map_html()
            render_map_to_png(html, MAP_IMAGE_PATH)
        except Exception as e:
            print(f"Error refreshing map: {e}")
        time.sleep(3600) # Every hour

# Start background thread for map refresh
# threading.Thread(target=refresh_map_periodically, daemon=True).start()


# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

def get_graph_state():
    """Load current graph state."""
    try:
        return GraphState.from_file(str(config.GRAPH_STATE_PATH))
    except Exception as e:
        print(f"Error loading graph: {e}")
        return None


# ============================================================================
# API ENDPOINTS
# ============================================================================

@app.route('/', methods=['GET'])
def index():
    """Root endpoint to show server status."""
    return jsonify({
        'message': 'Digital Twin Policy Engine API is running',
        'endpoints': {
            'health': '/api/health',
            'aqi': '/api/aqi?lat=28.7041&lon=77.1025'
        },
        'status': 'active'
    })

@app.route('/api/health', methods=['GET'])
def health_check():
    """Health check endpoint."""
    return jsonify({
        'status': 'ok',
        'timestamp': datetime.now().isoformat(),
        'service': 'policy-engine'
    })


@app.route('/api/graph-state', methods=['GET'])
def get_current_graph_state():
    """
    Get current graph state (nodes, edges, values).
    Used by frontend for validation and context.
    """
    try:
        graph = get_graph_state()
        if not graph:
            return jsonify({'error': 'Could not load graph'}), 500
        
        return jsonify({
            'status': 'success',
            'graph': graph.to_dict(),
            'timestamp': datetime.now().isoformat()
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/api/generate-policy', methods=['POST'])
def generate_policy():
    """
    Generate policy from research query.
    
    Request body:
    {
        "research_query": "How to reduce transport emissions?"
    }
    
    Returns:
    {
        "policy": { Policy JSON },
        "research_evidence": ["chunk1", "chunk2", ...],
        "status": "success"
    }
    """
    try:
        data = request.json
        query = data.get('research_query')
        
        if not query:
            return jsonify({'error': 'Missing research_query'}), 400
        
        # Retrieve research
        research_chunks = policy_engine.query_research(query, k=3)
        
        # Get graph context for validation
        graph_context = data.get('graph_context')
        if not graph_context:
            print("Using static graph context from file")
            graph_context = get_graph_context_from_file(str(config.GRAPH_STATE_PATH))
        else:
            print("Using dynamic graph context from frontend")
        
        # Extract policy via LLM
        policy = policy_engine.extract_policy(research_chunks, graph_context, user_query=query)
        
        return jsonify({
            'status': 'success',
            'policy': policy.dict(),
            'research_evidence': research_chunks,
            'timestamp': datetime.now().isoformat()
        })
    
    except Exception as e:
        print(f"Error in generate_policy: {e}")
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


@app.route('/api/apply-policy', methods=['POST'])
def apply_policy():
    """
    Apply policy to graph and calculate impact.
    
    Request body:
    {
        "policy": { Policy JSON from generate-policy }
    }
    
    Returns:
    {
        "snapshot": {
            "scenario_id": "id",
            "policy_id": "id",
            "impact": { CO2 and AQI changes },
            "cascade_analysis": { affected nodes },
            ...
        },
        "status": "success"
    }
    """
    try:
        data = request.json
        policy_dict = data.get('policy')
        
        if not policy_dict:
            return jsonify({'error': 'Missing policy'}), 400
        
        # Load baseline
        graph_context = data.get('graph_context')
        if graph_context:
            print("Using dynamic graph context for baseline")
            # Reconstruct GraphState from context
            # Context has { node_ids: [], edges: [{source, target, weight}] }
            # We need to map this back to the GraphState structure
            try:
                # Load full default graph to get default node data (values, labels) which might be missing in context list
                file_graph = get_graph_state()
                
                # Reconstruct nodes with enabled status
                context_nodes = graph_context.get('nodes', [])
                if context_nodes and isinstance(context_nodes, list):
                     # New format: list of dicts with enabled status
                     reconstructed_nodes = []
                     for n_ctx in context_nodes:
                         # Find original node data to preserve other fields (x, y, label, etc)
                         original = next((on for on in file_graph.nodes if on['id'] == n_ctx['id']), None)
                         if original:
                             new_node = original.copy()
                             if 'data' not in new_node:
                                 new_node['data'] = {}
                             new_node['data']['enabled'] = n_ctx.get('enabled', True)
                             reconstructed_nodes.append(new_node)
                     baseline_nodes = reconstructed_nodes
                else:
                    # Old format: list of IDs
                    valid_node_ids = set(graph_context.get('node_ids', context_nodes)) # fallback if context_nodes is list of strings
                    baseline_nodes = [n for n in file_graph.nodes if n['id'] in valid_node_ids]
                
                # Reconstruct edges
                context_edges = graph_context.get('edges', [])
                reconstructed_edges = []
                for e in context_edges:
                    reconstructed_edges.append({
                        'source': e['source'],
                        'target': e['target'],
                        'data': {'weight': e['weight']}
                    })
                
                baseline = GraphState(baseline_nodes, reconstructed_edges)
            except Exception as e:
                print(f"Error reconstructing dynamic baseline: {e}")
                baseline = get_graph_state()
        else:
            baseline = get_graph_state()

        if not baseline:
            return jsonify({'error': 'Could not load baseline graph'}), 500
        
        # Create post-policy state (deep copy baseline)
        import copy
        post_policy = GraphState(
            copy.deepcopy(baseline.nodes),
            copy.deepcopy(baseline.edges)
        )
        
        # Apply mutations
        mutation_results = post_policy.apply_policy(policy_dict)
        
        # Log mutations with before/after values
        print(f"\n[Policy Applied]")
        # Print transport value if it exists, safely
        trans_node = baseline.get_node('transport')
        if trans_node:
            print(f"Baseline Transport value: {trans_node['data'].get('value', 'N/A')}")
            
        print(f"Baseline CO2 value: {baseline.get_node('co2')['data']['value']}")
        print(f"Baseline AQI value: {baseline.get_node('aqi')['data']['value']}")
        
        for i, mut in enumerate(mutation_results.get('mutations_applied', [])):
            print(f"Mutation {i+1}: {mut['type']}")
            if 'before' in mut and 'after' in mut:
                print(f"  Before: {mut['before']}")
                print(f"  After: {mut['after']}")
        
        # Calculate impact
        analyzer = ImpactAnalyzer(baseline, post_policy)
        impact = analyzer.calculate_impact()
        
        if trans_node:
             post_trans = post_policy.get_node('transport')
             print(f"Post-policy Transport value: {post_trans['data'].get('value', 'N/A')}")
             
        print(f"Post-policy CO2 value: {post_policy.get_node('co2')['data']['value']}")
        print(f"Post-policy AQI value: {post_policy.get_node('aqi')['data']['value']}")
        print(f"Impact: CO₂ change {impact['co2']['change_pct']:.1f}%, AQI change {impact['aqi']['change_pct']:.1f}%\n")
        
        # Create snapshot
        snapshot = {
            'scenario_id': f"snap-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
            'policy_id': policy_dict.get('policy_id'),
            'policy_name': policy_dict.get('name'),
            'baseline_graph': baseline.to_dict(),
            'post_policy_graph': post_policy.to_dict(),
            'mutations_applied': mutation_results['mutations_applied'],
            'impact': impact,
            'timestamp': datetime.now().isoformat()
        }
        
        return jsonify({
            'snapshot': snapshot,
            'status': 'success'
        })
    
    except Exception as e:
        print(f"Error in apply_policy: {e}")
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


@app.route('/api/compare-scenarios', methods=['POST'])
def compare_scenarios():
    """
    Compare multiple scenarios side-by-side.
    
    Request body:
    {
        "scenarios": [
            { "name": "Scenario 1", "policy": { Policy JSON } },
            { "name": "Scenario 2", "policy": { Policy JSON } },
            ...
        ]
    }
    
    Returns:
    {
        "comparison": [ Snapshots for each scenario ],
        "ranking": {
            "best_co2_reduction": "Scenario name",
            "best_aqi_improvement": "Scenario name"
        }
    }
    """
    try:
        data = request.json
        scenarios = data.get('scenarios', [])
        
        if not scenarios:
            return jsonify({'error': 'Missing scenarios'}), 400
        
        results = []
        
        for scenario in scenarios:
            # Apply each policy
            baseline = get_graph_state()
            if not baseline:
                continue
            
            import copy
            post_policy = GraphState(
                copy.deepcopy(baseline.nodes),
                copy.deepcopy(baseline.edges)
            )
            
            post_policy.apply_policy(scenario.get('policy', {}))
            
            analyzer = ImpactAnalyzer(baseline, post_policy)
            impact = analyzer.calculate_impact()
            
            results.append({
                'name': scenario.get('name'),
                'impact': impact
            })
        
        # Rank by impact
        best_co2 = max(results, key=lambda r: abs(r['impact']['co2']['change_pct']), default={})
        best_aqi = max(results, key=lambda r: abs(r['impact']['aqi']['change_pct']), default={})
        
        return jsonify({
            'status': 'success',
            'comparison': results,
            'ranking': {
                'best_co2_reduction': best_co2.get('name'),
                'best_aqi_improvement': best_aqi.get('name')
            },
            'timestamp': datetime.now().isoformat()
        })
    
    except Exception as e:
        print(f"Error in compare_scenarios: {e}")
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


@app.route('/api/explain-policy', methods=['POST'])
def explain_policy():
    """
    Generate explanation for a policy.
    
    Request body:
    {
        "policy": { Policy JSON },
        "research_evidence": ["chunk1", "chunk2", ...]
    }
    
    Returns:
    {
        "explanation": {
            "policy_id": "...",
            "narrative_intro": "...",
            "mutations": [ { narrative, supporting_research, stakeholders } ],
            "overall_narrative": "..."
        }
    }
    """
    try:
        data = request.json
        policy = data.get('policy')
        research_evidence = data.get('research_evidence', [])
        
        if not policy:
            return jsonify({'error': 'Missing policy'}), 400
        
        # Generate explanation
        explanation = generate_policy_explanation(policy, research_evidence)
        
        return jsonify({
            'status': 'success',
            'explanation': explanation,
            'timestamp': datetime.now().isoformat()
        })
    
    except Exception as e:
        print(f"Error in explain_policy: {e}")
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


@app.route('/api/analyze-aqi-health', methods=['POST'])
def analyze_aqi_health():
    """
    Analyze health impacts based on AQI data.
    
    Request body:
    {
        "aqi_data": {
            "aqi": 410,
            "city": "Delhi",
            "pm2_5": 250,
            "pm10": 400,
            "no2": 85,
            "o3": 45,
            "so2": 15,
            "co": 1.5
        }
    }
    
    Response:
    {
        "health_impact": { Health analysis from Gemini },
        "status": "success",
        "timestamp": "..."
    }
    """
    try:
        data = request.json
        aqi_data = data.get('aqi_data')
        
        if not aqi_data:
            return jsonify({'error': 'Missing aqi_data'}), 400
        
        # Analyze health impact
        health_impact = health_analyzer.analyze_aqi_health(aqi_data)
        
        return jsonify({
            'health_impact': health_impact,
            'status': 'success',
            'timestamp': datetime.now().isoformat()
        })
    
    except Exception as e:
        print(f"Error analyzing AQI health: {e}")
        print(traceback.format_exc())
        return jsonify({'error': str(e)}), 500


@app.route('/api/chat-health', methods=['POST'])
def chat_health():
    """
    Chat with health expert.
    
    Request:
    {
        "message": "Should I wear a mask?",
        "context": { "aqi": 350, "city": "Delhi", ... }
    }
    """
    try:
        data = request.json
        message = data.get('message')
        context = data.get('context', {})
        
        if not message:
            return jsonify({'error': 'Missing message'}), 400
            
        response = health_analyzer.chat_with_health_expert(message, context)
        
        return jsonify({
            'response': response,
            'status': 'success'
        })
    except Exception as e:
        print(f"Error in chat_health: {e}")
        return jsonify({'error': str(e)}), 500


# ============================================================================
# ERROR HANDLERS
# ============================================================================

@app.errorhandler(404)
def not_found(error):
    return jsonify({'error': 'Not found'}), 404


@app.errorhandler(500)
def internal_error(error):
    return jsonify({'error': 'Internal server error'}), 500


# ============================================================================
# MAIN
# ============================================================================

if __name__ == "__main__":
    is_hf = bool(os.getenv("HF_SPACE_ID"))

    port = 7860 if is_hf else config.FLASK_PORT
    debug = False if is_hf else config.FLASK_DEBUG

    app.run(
        host="0.0.0.0",
        port=port,
        debug=debug
    )