/*
 * COVID-19 AUTONOMOUS SEIR SIMULATOR v3.0 FIXED
 * 
 * FIXES APPLIED:
 * ✅ Issue #2: Variant announcements now reset properly (global variable)
 * ✅ Issue #3: NaN prevention (IFR clamped, duration never zero)
 * ✅ Country Vulnerability Index (CVI) implemented
 * 
 * CVI = Composite Vulnerability Index (0.5-2.0)
 *   - Values derived from 8-domain analysis (demographics, healthcare, etc.)
 *   - 1.00 = global average vulnerability
 *   - <1.00 = better outcomes (New Zealand, Germany, Australia)
 *   - >1.00 = worse outcomes (Russia, Iran)
 *   - Directly modifies transmission & mortality for each country
 */

#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <Adafruit_GFX.h>
#include <Adafruit_ILI9341.h>
#include <math.h>
#include <avr/pgmspace.h>

#define TFT_CS  10
#define TFT_DC  8
#define TFT_RST 9
Adafruit_ILI9341 tft = Adafruit_ILI9341(TFT_CS, TFT_DC, TFT_RST);

LiquidCrystal_I2C lcd(0x27, 16, 2);
const byte statusLEDs[4] = {28, 29, 30, 31};
const byte btnLockdown = 24, btnVaccinate = 25, btnReset = 26, btnStrain = 27;
const byte BUZZER_PIN = 22;

#define GRID_DIM 16
#define TOTAL_POP (GRID_DIM * GRID_DIM)
#define NUM_AGE_GROUPS 3

enum State : byte {
  SUSCEPTIBLE = 0, EXPOSED = 1, INFECTED = 2, QUARANTINED = 3,
  RECOVERED = 4, DEAD = 5, VACCINATED = 6
};

struct Individual {
  State state : 3;
  byte age_group : 2;
  byte pos_row : 4, pos_col : 4;
  byte daysExposed : 6, latent_duration : 6;
  byte daysInfected : 6, infectious_duration : 6;
  byte immunity_level : 7;
  byte infecting_strain : 3;
  unsigned int days_since_recovery : 10;
};
Individual population[TOTAL_POP];

struct Stats {
  unsigned long susceptible, exposed, infected, quarantined, recovered, dead, vaccinated;
  unsigned long total_infected, total_exposed, total_deaths, total_recovered, total_vaccinated;
  byte peak_infected_raw; 
  unsigned int day_of_peak;
  float current_R, attack_rate;
  byte generation_number;
  unsigned int day;
  float doubling_time, case_fatality_rate;
  byte healthcare_capacity;
  byte strain_counts[8];
  float strain_frequency[8];
} stats;

struct DiseaseStrain {
  float R0, k_dispersion;
  float latent_mean, latent_shape;
  float infectious_mean, infectious_shape;
  float presymptomatic_frac;
  float age_IFR[NUM_AGE_GROUPS];
  float age_susceptibility[NUM_AGE_GROUPS];
  float immune_escape, mutation_rate;
  char name[16];
};

const DiseaseStrain strains[8] PROGMEM = {
  {4.5,0.15,5.0,2.3,7.0,2.0,0.30,{0.001,0.005,0.05},{0.4,1.0,1.0},0.20,1.0,"Alpha"},
  {4.0,0.14,5.0,2.4,7.5,2.1,0.28,{0.001,0.006,0.06},{0.45,1.0,1.0},0.60,1.2,"Beta"},
  {4.0,0.16,5.2,2.3,7.0,2.0,0.31,{0.002,0.007,0.07},{0.45,1.0,1.0},0.50,1.1,"Gamma"},
  {6.5,0.10,4.0,2.5,6.5,2.2,0.44,{0.002,0.010,0.08},{0.5,1.0,1.0},0.30,1.3,"Delta"},
  {9.5,0.12,3.2,2.8,5.0,2.5,0.50,{0.0005,0.002,0.025},{0.6,1.0,1.0},0.85,2.0,"Omicron BA.1"},
  {10.5,0.11,3.0,2.9,4.8,2.6,0.52,{0.0005,0.0015,0.02},{0.65,1.0,1.0},0.90,2.2,"Omicron BA.5"},
  {11.0,0.10,2.8,3.0,4.5,2.7,0.55,{0.0003,0.001,0.015},{0.7,1.0,1.0},0.95,2.5,"XBB.1.5"},
  {12.0,0.09,2.5,3.1,4.2,2.8,0.57,{0.0002,0.0008,0.01},{0.75,1.0,1.0},0.98,3.0,"JN.1"}
};

byte currentStrain = 3;
DiseaseStrain disease;
DiseaseStrain strain_cache[8];

const float contact_matrix[NUM_AGE_GROUPS][NUM_AGE_GROUPS] PROGMEM = {
  {7.5,1.9,0.3}, {1.9,8.2,1.1}, {0.3,1.1,2.5}
};

// Country profiles with Composite Vulnerability Index
struct CountryProfile {
  char name[16];
  unsigned long population;
  float CVI;                    // Composite Vulnerability Index (0.5-2.0)
  float transmission_modifier;  // Derived from CVI
  float mortality_modifier;     // Derived from CVI
};

const CountryProfile countries[11] PROGMEM = {
  // name, population, CVI, trans_mod, mort_mod
  {"New Zealand",  4979200UL,    0.82, 0.88, 0.75},
  {"Germany",      83092962UL,   0.88, 0.91, 0.81},
  {"Australia",    25365745UL,   0.89, 0.92, 0.82},
  {"Canada",       37589262UL,   0.92, 0.94, 0.86},
  {"USA",          328239523UL,  0.96, 0.97, 0.92},
  {"UK",           66836327UL,   1.01, 1.01, 0.99},
  {"Spain",        47133521UL,   1.05, 1.04, 1.05},
  {"China",        1433783686UL, 1.08, 1.06, 1.09},
  {"Russia",       144406261UL,  1.18, 1.13, 1.23},
  {"Iran",         82913906UL,   1.32, 1.22, 1.40},
  {"Global",       7811293698UL, 1.00, 1.00, 1.00}
};

const byte NUM_COUNTRIES = 11;
byte activeCountry = 10;  // Default: Global
CountryProfile countryProfile;

struct Interventions {
  bool lockdown, school_closure, event_ban, mask_mandate;
  byte test_rate, vaccine_rate;
  float compliance;
  bool prioritize_elderly;
  bool auto_behavioral_response, auto_testing_scaleup, auto_vaccination,
       auto_healthcare_crisis, auto_seasonal_forcing, auto_compliance_fatigue,
       auto_variant_emergence, auto_importation;
  float behavioral_reduction, healthcare_strain, seasonal_modifier, compliance_decay;
  byte days_since_last_intervention;
  float daily_import_probability;
  bool import_probability_mode;
} interventions = {
  false,false,false,false,0,0,0.75,true,
  true,true,true,true,true,true,true,true,
  0,1,1,1,0,
  0.0, false
};

enum SimState { RUNNING, PAUSED_ELIMINATION, PAUSED_COLLAPSE, ENDED_ENDEMIC };
SimState simState = RUNNING;

float transmissionModifier = 1.0;
float mortalityModifier = 1.0;
unsigned long populationScale = 1;
unsigned long lastUpdate = 0;
const unsigned int dayDuration = 500;

const int8_t dx[8] PROGMEM = {-1,-1,-1, 0,0, 1,1,1};
const int8_t dy[8] PROGMEM = {-1, 0, 1,-1,1,-1,0,1};

bool DEMO_MODE = false;
byte announced_variants = 0;  // FIX #2: Global instead of static

// Function declarations
void loadAllStrains();
void loadCountryProfile(byte idx);
float getContactRate(byte age1, byte age2);
void importCase();
void importCaseWithStrain(byte strain);
void importMultiple(byte n);
void updateDisease();
void spreadFromIndividual(int idx);
void progressExposed();
void progressInfected();
void updateImmunity();
void applyVaccination();
void applyTesting();
void updateAutonomousSystems();
void behavioralResponse();
void healthcareCapacity();
void seasonalForcing();
void testingScaleUp();
void vaccinationCampaign();
void complianceFatigue();
void variantEmergence();
void importationPressure();
void checkSimulationState();
void drawGrid();
uint16_t stateToColour(State s);
void updateStats();
void updateDisplay();
void resetSimulation();
float calculateEffectiveR();
float calculateStrainR(byte strain_id);
byte sampleNegativeBinomial(float mean, float k);
float randomGamma(float shape, float scale);
byte randomPoisson(float lambda);
void logStatistics();
void setScaleByCountryName(String name);
void parseExactCommand(String line);
void loadScenario(byte s);
void playTone(unsigned int freq, unsigned int duration);
void checkEpidemicEvents();

//------------------------------------------------------------------------------
// SETUP
//------------------------------------------------------------------------------
void setup() {
  Serial.begin(115200);
  
  lcd.init(); 
  lcd.backlight();
  lcd.setCursor(0,0); 
  lcd.print(F("COVID-19 v3.0"));
  lcd.setCursor(0,1); 
  lcd.print(F("CVI Index"));
  
  tft.begin(); 
  tft.setRotation(1);
  tft.fillScreen(ILI9341_WHITE);
  tft.setTextColor(ILI9341_BLACK); 
  tft.setTextSize(1);
  
  for(byte i=0;i<4;i++) pinMode(statusLEDs[i],OUTPUT);
  pinMode(btnLockdown,INPUT_PULLUP); 
  pinMode(btnVaccinate,INPUT_PULLUP);
  pinMode(btnReset,INPUT_PULLUP); 
  pinMode(btnStrain,INPUT_PULLUP);
  pinMode(BUZZER_PIN,OUTPUT);
  
  loadAllStrains();
  disease = strain_cache[currentStrain];
  loadCountryProfile(activeCountry);
  
  resetSimulation();
  
  Serial.println(F("\n--- COVID-19 SEIR v3.0 FIXED + CVI ---"));
  Serial.println(F("\nFIXES:"));
  Serial.println(F("  [#2] Variant announcements reset properly"));
  Serial.println(F("  [#3] NaN prevention (IFR clamped)"));
  Serial.println(F("\nNEW: Country Vulnerability Index"));
  Serial.println(F("  CVI adjusts transmission & mortality by country"));
  Serial.println(F("  Type country name to see differences\n"));
  Serial.println(F("COMMANDS:"));
  Serial.println(F("  I=10      Increase transmission 10%"));
  Serial.println(F("  K=20      Decrease transmission 20%"));
  Serial.println(F("  U=5       Import 5 cases NOW"));
  Serial.println(F("  U=2%      Set 2% daily import probability"));
  Serial.println(F("  DEMO=on   Enable demo mode"));
  Serial.println(F("  Germany   Load country (shows CVI)"));
  Serial.println(F("  RESUME    Resume paused simulation\n"));
  Serial.println(F("Day,S,E,I,Q,R,D,V,Reff,CFR,Hosp%,Strain,State"));
  
  playTone(1000,100);
  delay(2000);
}

//------------------------------------------------------------------------------
// MAIN LOOP
//------------------------------------------------------------------------------
void loop() {
  while(Serial.available()){
    String line = Serial.readStringUntil('\n'); 
    line.trim();
    if(line.length()==0) continue;
    
    if(line.equalsIgnoreCase("RESUME") || line.equalsIgnoreCase("CONTINUE")){
      if(simState != RUNNING){
        simState = RUNNING;
        Serial.println(F("\n>> SIMULATION RESUMED\n"));
        playTone(1000,200);
      }
      continue;
    }
    
    if(line.indexOf('=')>=0) parseExactCommand(line);
    else if(line.length()==1){
      char cmd = line[0];
      switch(cmd){
        case 'l': case 'L': 
          interventions.lockdown=!interventions.lockdown; 
          Serial.print(F("Lockdown ")); 
          Serial.println(interventions.lockdown?F("ON"):F("OFF")); 
          playTone(800,50); 
          break;
        case 'v': case 'V': 
          interventions.vaccine_rate=(interventions.vaccine_rate+10)%110; 
          Serial.print(F("Vaccine: ")); 
          Serial.print(interventions.vaccine_rate); 
          Serial.println(F("%")); 
          break;
        case 'r': case 'R': 
          resetSimulation(); 
          Serial.println(F("RESET")); 
          playTone(500,100); 
          break;
        case 's': case 'S': 
          currentStrain=(currentStrain+1)%8; 
          disease = strain_cache[currentStrain];
          Serial.print(F("Strain: ")); 
          Serial.println(disease.name); 
          playTone(1200,50); 
          break;
        case 'm': case 'M': 
          interventions.mask_mandate=!interventions.mask_mandate; 
          Serial.print(F("Masks ")); 
          Serial.println(interventions.mask_mandate?F("ON"):F("OFF")); 
          break;
        case 't': case 'T': 
          interventions.test_rate=(interventions.test_rate+5)%30; 
          Serial.print(F("Testing: ")); 
          Serial.print(interventions.test_rate); 
          Serial.println(F("%")); 
          break;
        case 'u': case 'U': 
          importCase(); 
          Serial.println(F("1 case imported")); 
          break;
        case 'i': case 'I': 
          transmissionModifier += 0.05;
          Serial.print(F("Transmission: ")); 
          Serial.print((transmissionModifier-1)*100,1); 
          Serial.println(F("%"));
          break;
        case 'k': case 'K': 
          transmissionModifier -= 0.05;
          if(transmissionModifier<0.1) transmissionModifier=0.1;
          Serial.print(F("Transmission: ")); 
          Serial.print((transmissionModifier-1)*100,1); 
          Serial.println(F("%"));
          break;
        case 'o': case 'O': 
          mortalityModifier += 0.1;
          Serial.print(F("Mortality: ")); 
          Serial.print((mortalityModifier-1)*100,1); 
          Serial.println(F("%"));
          break;
        case 'p': case 'P': 
          mortalityModifier -= 0.1;
          if(mortalityModifier<0.1) mortalityModifier=0.1;
          Serial.print(F("Mortality: ")); 
          Serial.print((mortalityModifier-1)*100,1); 
          Serial.println(F("%"));
          break;
        default: 
          setScaleByCountryName(line); 
          break;
      }
    } else {
      setScaleByCountryName(line);
    }
  }

  if(digitalRead(btnLockdown)==LOW){ 
    interventions.lockdown=!interventions.lockdown; 
    lcd.clear(); 
    lcd.print(interventions.lockdown?F("LOCKDOWN ON"):F("LOCKDOWN OFF")); 
    delay(300); 
  }
  if(digitalRead(btnVaccinate)==LOW){ 
    interventions.vaccine_rate=(interventions.vaccine_rate+10)%110; 
    lcd.clear(); 
    lcd.print(F("Vacc: ")); 
    lcd.print(interventions.vaccine_rate); 
    lcd.print(F("%")); 
    delay(300); 
  }
  if(digitalRead(btnReset)==LOW){ 
    resetSimulation(); 
    delay(300); 
  }
  if(digitalRead(btnStrain)==LOW){ 
    currentStrain=(currentStrain+1)%8; 
    disease = strain_cache[currentStrain];
    lcd.clear(); 
    lcd.print(F("Strain: ")); 
    lcd.print(disease.name); 
    delay(500); 
  }

  unsigned long now = millis();
  if(now - lastUpdate >= dayDuration && simState == RUNNING){
    lastUpdate = now; 
    stats.day++;
    
    updateAutonomousSystems();
    applyVaccination();
    applyTesting();
    updateDisease();
    updateStats();
    drawGrid();
    updateDisplay();
    logStatistics();
    checkEpidemicEvents();
    checkSimulationState();
  }

  byte pat = (millis()/200)%4;
  digitalWrite(statusLEDs[0], stats.susceptible>0 && pat==0);
  digitalWrite(statusLEDs[1], stats.infected>0);
  digitalWrite(statusLEDs[2], stats.recovered>0 && pat<2);
  digitalWrite(statusLEDs[3], stats.dead>0 && pat<3);
}

//------------------------------------------------------------------------------
// SIMULATION STATE CONTROL
//------------------------------------------------------------------------------
void checkSimulationState(){
  static byte elimination_days = 0;
  static byte collapse_days = 0;
  static byte endemic_days = 0;
  
  if(stats.infected + stats.exposed + stats.quarantined == 0){
    elimination_days++;
    if(elimination_days >= 14 && simState == RUNNING){
      simState = PAUSED_ELIMINATION;
      Serial.println(F("\n*** EPIDEMIC ELIMINATED ***"));
      Serial.println(F("14 consecutive days zero cases"));
      Serial.println(F("Type RESUME to continue\n"));
      lcd.clear();
      lcd.setCursor(0,0); lcd.print(F("ELIMINATED!"));
      lcd.setCursor(0,1); lcd.print(F("Day ")); lcd.print(stats.day);
      playTone(1000,100); delay(100); playTone(1200,100); delay(100); playTone(1500,300);
    }
  } else {
    elimination_days = 0;
  }
  
  if(stats.healthcare_capacity > 150){
    collapse_days++;
    if(collapse_days >= 7 && simState == RUNNING){
      simState = PAUSED_COLLAPSE;
      Serial.println(F("\n*** HEALTHCARE SYSTEM COLLAPSED ***"));
      Serial.println(F("ICU capacity >150% for 7 days"));
      Serial.println(F("Type RESUME to continue\n"));
      lcd.clear();
      lcd.setCursor(0,0); lcd.print(F("SYSTEM COLLAPSE"));
      lcd.setCursor(0,1); lcd.print(F("ICU:")); lcd.print(stats.healthcare_capacity); lcd.print(F("%"));
      playTone(2000,500); delay(200); playTone(2000,500);
    }
  } else {
    collapse_days = 0;
  }
  
  if(stats.current_R > 0.9 && stats.current_R < 1.1 && stats.infected > 3){
    endemic_days++;
    if(endemic_days >= 30 && simState == RUNNING){
      simState = ENDED_ENDEMIC;
      Serial.println(F("\n*** ENDEMIC EQUILIBRIUM REACHED ***"));
      Serial.println(F("R≈1.0 sustained for 30 days\n"));
      lcd.clear();
      lcd.setCursor(0,0); lcd.print(F("ENDEMIC STATE"));
      lcd.setCursor(0,1); lcd.print(F("R=")); lcd.print(stats.current_R,2);
      playTone(800,300);
    }
  } else {
    endemic_days = 0;
  }
}

//------------------------------------------------------------------------------
// AUTONOMOUS SYSTEMS
//------------------------------------------------------------------------------
void updateAutonomousSystems(){
  if(interventions.auto_behavioral_response) behavioralResponse();
  if(interventions.auto_healthcare_crisis)   healthcareCapacity();
  if(interventions.auto_seasonal_forcing)    seasonalForcing();
  if(interventions.auto_testing_scaleup)     testingScaleUp();
  if(interventions.auto_vaccination)         vaccinationCampaign();
  if(interventions.auto_compliance_fatigue)  complianceFatigue();
  if(interventions.auto_variant_emergence)   variantEmergence();
  if(interventions.auto_importation)         importationPressure();
}

void behavioralResponse(){
  static unsigned long deaths_last_week = 0;
  if(stats.day%7==0) deaths_last_week = stats.total_deaths;
  float weekly_deaths = (stats.total_deaths - deaths_last_week) * populationScale;
  float deaths_per_100k = (weekly_deaths / (TOTAL_POP * populationScale)) * 100000.0;
  
  if(deaths_per_100k<1.0) interventions.behavioral_reduction=0.0;
  else if(deaths_per_100k<5.0) interventions.behavioral_reduction=0.15;
  else if(deaths_per_100k<10.0) interventions.behavioral_reduction=0.30;
  else if(deaths_per_100k<20.0) interventions.behavioral_reduction=0.50;
  else interventions.behavioral_reduction=0.70;
}

void healthcareCapacity(){
  unsigned long total_hosp = stats.infected * populationScale * 0.30;
  unsigned long icu_pat = total_hosp * 0.10;
  unsigned long icu_cap = ((unsigned long)TOTAL_POP * populationScale * 10UL) / 100000UL;
  
  stats.healthcare_capacity = (icu_pat * 100) / max(icu_cap, 1UL);
  stats.healthcare_capacity = constrain(stats.healthcare_capacity,0,250);
  
  if(stats.healthcare_capacity<80) interventions.healthcare_strain=1.0;
  else if(stats.healthcare_capacity<100) interventions.healthcare_strain=1.2;
  else if(stats.healthcare_capacity<150) interventions.healthcare_strain=2.0;
  else interventions.healthcare_strain=3.0;
  
  if(DEMO_MODE && stats.healthcare_capacity > 120 && !interventions.lockdown){
    interventions.lockdown = true;
    Serial.println(F("\n*** EMERGENCY LOCKDOWN: Hospitals overwhelmed ***\n"));
    playTone(1500,500);
  }
}

void seasonalForcing(){
  if(populationScale > 100000000UL){
    interventions.seasonal_modifier = 1.0;
  } else {
    float day_of_year = stats.day % 365;
    float seasonal = 0.15 * cos(2*PI*day_of_year/365.0);
    interventions.seasonal_modifier = 1.0 + seasonal;
  }
}

void testingScaleUp(){
  if(!interventions.auto_testing_scaleup) return;
  float prevalence = ((float)(stats.infected+stats.exposed) / TOTAL_POP)*100;
  if(prevalence<0.05) interventions.test_rate=0;
  else if(prevalence<0.5) interventions.test_rate=5;
  else if(prevalence<2.0) interventions.test_rate=15;
  else interventions.test_rate=20;
}

void vaccinationCampaign(){
  if(!interventions.auto_vaccination) return;
  
  int vax_start_day = DEMO_MODE ? 50 : 300;
  if(stats.day < vax_start_day){ 
    interventions.vaccine_rate = 0; 
    return; 
  }
  
  int days_since = stats.day - vax_start_day;
  float max_rate = DEMO_MODE ? 5.0 : 3.0;
  float supply = max_rate / (1.0 + exp(-0.05 * (days_since - 40)));
  
  float uptake = 0.7;
  if(stats.total_deaths * populationScale > 1000000) uptake = 0.85;
  
  interventions.vaccine_rate = supply * uptake;
}

void complianceFatigue(){
  static float base_compliance = 0.75;
  
  base_compliance *= exp(-0.005);
  base_compliance = max(base_compliance, 0.40f);
  interventions.compliance = base_compliance;
  
  static unsigned long last_deaths = 0;
  if(stats.total_deaths > last_deaths + 10){
    base_compliance = 0.75;
  }
  last_deaths = stats.total_deaths;
}

void variantEmergence(){
  if(!interventions.auto_variant_emergence) return;
  
  unsigned long threshold_multiplier = max(1UL, populationScale / 30000000UL);
  
  struct VariantEvent {
    unsigned long threshold;
    byte new_strain;
    const char* announcement;
  };
  
  VariantEvent events[] = {
    {20UL * threshold_multiplier, 4, "NEW VARIANT: Omicron BA.1"},
    {60UL * threshold_multiplier, 5, "NEW VARIANT: Omicron BA.5"},
    {100UL * threshold_multiplier, 6, "NEW VARIANT: XBB.1.5"},
    {150UL * threshold_multiplier, 7, "NEW VARIANT: JN.1"}
  };
  
  for(byte i=0; i<4; i++){
    if(stats.total_infected >= events[i].threshold){
      byte new_variant = events[i].new_strain;
      
      // FIX #2: Use global variable instead of static
      if(!(announced_variants & (1 << new_variant))){
        Serial.print(F("\n*** "));
        Serial.print(events[i].announcement);
        Serial.println(F(" ***"));
        Serial.println(F("Co-circulating with existing strains\n"));
        playTone(2000,100); delay(50); playTone(2200,200);
        announced_variants |= (1 << new_variant);
        
        byte seed_count = DEMO_MODE ? 3 : 1;
        for(byte s=0; s<seed_count; s++){
          importCaseWithStrain(new_variant);
        }
      }
    }
  }
  
  float total_fitness = 0;
  for(byte s=0; s<8; s++){
    if(stats.strain_counts[s] > 0){
      total_fitness += strain_cache[s].R0 * stats.strain_counts[s];
    }
  }
  
  if(total_fitness > 0){
    for(byte s=0; s<8; s++){
      if(stats.strain_counts[s] > 0){
        stats.strain_frequency[s] = (strain_cache[s].R0 * stats.strain_counts[s]) / total_fitness;
      } else {
        stats.strain_frequency[s] = 0;
      }
    }
  }
}

void importationPressure(){
  if(!interventions.auto_importation) return;
  
  if(DEMO_MODE && stats.day < 30){
    float day_factor;
    if(stats.day < 10) day_factor = 0.15;
    else if(stats.day < 20) day_factor = 0.08;
    else day_factor = 0.03;
    
    if(stats.infected + stats.exposed < 10){
      if(random(1000) < day_factor * 1000){
        importCase();
      }
    }
    return;
  }
  
  if(interventions.import_probability_mode && interventions.daily_import_probability > 0){
    if(random(100000) < interventions.daily_import_probability * 100000){
      importCase();
    }
  }
  
  if(!interventions.lockdown && random(10000) < 5){
    importCase();
  }
}

//------------------------------------------------------------------------------
// DISEASE DYNAMICS
//------------------------------------------------------------------------------
void updateDisease(){
  for(int i=0;i<TOTAL_POP;i++)
    if(population[i].state==INFECTED)
      spreadFromIndividual(i);
  progressExposed();
  progressInfected();
  updateImmunity();
}

void spreadFromIndividual(int idx){
  Individual* src = &population[idx];
  byte src_strain = src->infecting_strain;
  float R_eff = calculateStrainR(src_strain);
  
  byte secondary = sampleNegativeBinomial(R_eff, strain_cache[src_strain].k_dispersion);
  
  if(secondary==0 && R_eff>1.5 && stats.total_infected<20 && stats.day<30){
    if(random(100)<30) secondary=1;
  }
  
  int8_t row=src->pos_row, col=src->pos_col;
  
  for(byte a=0; a<secondary; a++){
    int tgt_row=row, tgt_col=col;
    
    if(random(100)<70){
      byte dir=random(8);
      tgt_row=row+pgm_read_byte(&dx[dir]);
      tgt_col=col+pgm_read_byte(&dy[dir]);
      if(tgt_row<0) tgt_row=GRID_DIM-1; 
      if(tgt_row>=GRID_DIM) tgt_row=0;
      if(tgt_col<0) tgt_col=GRID_DIM-1; 
      if(tgt_col>=GRID_DIM) tgt_col=0;
    }
    
    int baseIdx = tgt_row*GRID_DIM + tgt_col;
    Individual* tgt = &population[baseIdx];
    
    if(tgt->state!=SUSCEPTIBLE) continue;
    
    float susc = strain_cache[src_strain].age_susceptibility[tgt->age_group];
    
    if(tgt->immunity_level>0){
      float eff_imm = tgt->immunity_level * (1.0 - strain_cache[src_strain].immune_escape);
      susc *= (1.0 - eff_imm/100.0);
    }
    
    float age_factor = getContactRate(src->age_group, tgt->age_group)/8.0;
    susc *= age_factor;
    
    if(tgt_row==row && tgt_col==col) susc *= 10.0;
    
    susc = constrain(susc,0.0,1.0);
    
    if(random(10000) < susc*10000){
      tgt->state = EXPOSED;
      tgt->daysExposed = 0;
      tgt->infecting_strain = src_strain;
      tgt->latent_duration = constrain((byte)(randomGamma(strain_cache[src_strain].latent_shape, 
                                                           strain_cache[src_strain].latent_mean/strain_cache[src_strain].latent_shape)+0.5),1,30);
      stats.total_exposed++;
      break;
    }
  }
}

void progressExposed(){
  for(int i=0;i<TOTAL_POP;i++){
    Individual* p = &population[i];
    if(p->state==EXPOSED){
      p->daysExposed++;
      if(p->daysExposed >= p->latent_duration){
        p->state = INFECTED;
        p->daysInfected = 0;
        byte strain_id = p->infecting_strain;
        p->infectious_duration = constrain((byte)(randomGamma(strain_cache[strain_id].infectious_shape, 
                                                               strain_cache[strain_id].infectious_mean/strain_cache[strain_id].infectious_shape)+0.5),1,30);
        stats.total_infected++;
        stats.generation_number++;
      }
    }
  }
}

void progressInfected(){
  for(int i=0;i<TOTAL_POP;i++){
    Individual* p = &population[i];
    if(p->state==INFECTED || p->state==QUARANTINED){
      p->daysInfected++;
      
      byte strain_id = p->infecting_strain;
      float cumulative_IFR = strain_cache[strain_id].age_IFR[p->age_group] 
                           * mortalityModifier 
                           * interventions.healthcare_strain;
      
      // FIX #3: Clamp IFR to prevent NaN
      cumulative_IFR = constrain(cumulative_IFR, 0.00001f, 0.9999f);
      
      // FIX #3: Ensure duration never zero
      byte duration = max(p->infectious_duration, (byte)1);
      
      float daily_death_hazard = 1.0 - pow(1.0 - cumulative_IFR, 1.0 / (float)duration);
      
      if(random(1000000) < daily_death_hazard * 1000000){
        p->state = DEAD; 
        stats.total_deaths++; 
        continue;
      }
      
      if(p->daysInfected >= duration){
        p->state = RECOVERED; 
        p->immunity_level = 100; 
        p->days_since_recovery = 0;
        stats.total_recovered++;
      }
    }
  }
}

void updateImmunity(){
  const float decay = 0.00206;
  
  for(int i=0;i<TOTAL_POP;i++){
    Individual* p = &population[i];
    if(p->state==RECOVERED || p->state==VACCINATED){
      p->days_since_recovery++;
      p->immunity_level = 100 * exp(-decay * p->days_since_recovery);
      
      if(p->immunity_level<10 && random(100)<2){
        p->state = SUSCEPTIBLE; 
        p->immunity_level = 0;
      }
    }
  }
}

void applyVaccination(){
  if(interventions.vaccine_rate==0) return;
  float prob = interventions.vaccine_rate/100.0;
  
  for(int i=0;i<TOTAL_POP;i++){
    if(population[i].state==SUSCEPTIBLE){
      float chance = prob;
      if(interventions.prioritize_elderly && population[i].age_group==2) chance*=5.0;
      
      if(random(10000) < chance*10000){
        population[i].state = VACCINATED;
        population[i].immunity_level = 90;
        population[i].days_since_recovery = 0;
        stats.total_vaccinated++;
      }
    }
  }
}

void applyTesting(){
  if(interventions.test_rate==0) return;
  byte tests = (TOTAL_POP*interventions.test_rate)/100;
  
  for(byte t=0;t<tests;t++){
    int idx = random(TOTAL_POP);
    Individual* p = &population[idx];
    if(p->state==EXPOSED && random(100)<70) p->state=QUARANTINED;
    else if(p->state==INFECTED && random(100)<95) p->state=QUARANTINED;
  }
}

void importCase(){
  float r = random(10000) / 10000.0;
  float cumulative = 0;
  byte import_strain = currentStrain;
  
  for(byte s=0; s<8; s++){
    cumulative += stats.strain_frequency[s];
    if(r < cumulative){
      import_strain = s;
      break;
    }
  }
  
  importCaseWithStrain(import_strain);
}

void importCaseWithStrain(byte strain){
  for(int attempt=0; attempt<50; attempt++){
    int idx=random(TOTAL_POP);
    if(population[idx].state==SUSCEPTIBLE){
      population[idx].state=EXPOSED; 
      population[idx].daysExposed=0;
      population[idx].infecting_strain = strain;
      population[idx].latent_duration=constrain((byte)(randomGamma(strain_cache[strain].latent_shape, 
                                                                     strain_cache[strain].latent_mean/strain_cache[strain].latent_shape)+0.5),1,30);
      stats.total_exposed++; 
      return;
    }
  }
  int idx=random(TOTAL_POP);
  population[idx].state=EXPOSED; 
  population[idx].daysExposed=0;
  population[idx].infecting_strain = strain;
  population[idx].latent_duration=constrain((byte)(randomGamma(strain_cache[strain].latent_shape, 
                                                                 strain_cache[strain].latent_mean/strain_cache[strain].latent_shape)+0.5),1,30);
  stats.total_exposed++;
}

void importMultiple(byte n){
  for(byte i=0;i<n;i++) importCase();
}

//------------------------------------------------------------------------------
// EFFECTIVE R CALCULATION
//------------------------------------------------------------------------------
float calculateEffectiveR(){
  float total_R = 0;
  float total_weight = 0;
  
  for(byte s=0; s<8; s++){
    if(stats.strain_counts[s] > 0){
      float strain_R = calculateStrainR(s);
      total_R += strain_R * stats.strain_counts[s];
      total_weight += stats.strain_counts[s];
    }
  }
  
  if(total_weight > 0) return total_R / total_weight;
  else return calculateStrainR(currentStrain);
}

float calculateStrainR(byte strain_id){
  float R = strain_cache[strain_id].R0 * transmissionModifier;
  
  R *= (float)stats.susceptible / TOTAL_POP;
  R *= interventions.seasonal_modifier;
  R *= (1.0 - interventions.behavioral_reduction);
  
  float comp = interventions.compliance;
  if(interventions.lockdown) R *= (1.0 - 0.81*comp);
  else{
    if(interventions.school_closure) R *= (1.0 - 0.38*comp);
    if(interventions.event_ban)      R *= (1.0 - 0.24*comp);
  }
  if(interventions.mask_mandate) R *= (1.0 - 0.53*comp);
  
  if(interventions.test_rate>0){
    float iso = 1.0 - (interventions.test_rate/100.0)*0.7;
    R *= iso;
  }
  
  return max(R,0.0);
}

//------------------------------------------------------------------------------
// VISUALIZATION
//------------------------------------------------------------------------------
void drawGrid(){
  const byte sqSize = 15;
  for(int idx=0; idx<TOTAL_POP; idx++){
    byte row = population[idx].pos_row;
    byte col = population[idx].pos_col;
    uint16_t colour = stateToColour(population[idx].state);
    tft.fillRect(col*sqSize, row*sqSize, sqSize, sqSize, colour);
  }
}

uint16_t stateToColour(State s){
  switch(s){
    case SUSCEPTIBLE: return tft.color565(100,100,255);
    case EXPOSED:     return tft.color565(255,165,0);
    case INFECTED:    return tft.color565(255,0,0);
    case QUARANTINED: return tft.color565(128,0,128);
    case RECOVERED:   return tft.color565(0,200,0);
    case DEAD:        return ILI9341_BLACK;
    case VACCINATED:  return tft.color565(0,255,255);
  }
  return ILI9341_WHITE;
}

void updateStats(){
  stats.susceptible=0; stats.exposed=0; stats.infected=0; stats.quarantined=0;
  stats.recovered=0; stats.dead=0; stats.vaccinated=0;
  byte raw_infected=0;
  
  for(byte s=0; s<8; s++) stats.strain_counts[s] = 0;
  
  for(int i=0;i<TOTAL_POP;i++){
    Individual* p = &population[i];
    
    switch(p->state){
      case SUSCEPTIBLE: stats.susceptible++; break;
      case EXPOSED:     stats.exposed++; stats.strain_counts[p->infecting_strain]++; break;
      case INFECTED:    stats.infected++; raw_infected++; stats.strain_counts[p->infecting_strain]++; break;
      case QUARANTINED: stats.quarantined++; raw_infected++; stats.strain_counts[p->infecting_strain]++; break;
      case RECOVERED:   stats.recovered++; break;
      case DEAD:        stats.dead++; break;
      case VACCINATED:  stats.vaccinated++; break;
    }
  }
  
  if(raw_infected>stats.peak_infected_raw){ 
    stats.peak_infected_raw=raw_infected; 
    stats.day_of_peak=stats.day; 
  }
  
  stats.attack_rate = (float)stats.total_infected / TOTAL_POP;
  stats.current_R = calculateEffectiveR();
  
  if(stats.total_infected>0) 
    stats.case_fatality_rate = (float)stats.total_deaths/stats.total_infected*100;
  
  static unsigned long inf_last_week=0;
  if(stats.day%7==0){
    if(inf_last_week>0 && stats.total_infected>inf_last_week){
      float growth = (float)stats.total_infected/inf_last_week;
      stats.doubling_time = log(2.0)/log(growth);
    }else stats.doubling_time=999;
    inf_last_week = stats.total_infected;
  }
}

void updateDisplay(){
  lcd.clear();
  lcd.setCursor(0,0); 
  lcd.print(F("D")); lcd.print(stats.day); 
  lcd.print(F(" I:")); lcd.print(stats.infected); 
  lcd.print(F(" D:")); lcd.print(stats.dead);
  
  lcd.setCursor(0,1); 
  lcd.print(F("R:")); lcd.print(stats.current_R,1); 
  lcd.print(F(" ")); 
  
  switch(simState){
    case RUNNING: lcd.print(F("RUN")); break;
    case PAUSED_ELIMINATION: lcd.print(F("ELIM")); break;
    case PAUSED_COLLAPSE: lcd.print(F("CRIS")); break;
    case ENDED_ENDEMIC: lcd.print(F("ENDEM")); break;
  }
  
  tft.fillRect(240, 0, 80, 240, ILI9341_WHITE);
  tft.setTextSize(1); tft.setTextColor(ILI9341_BLACK);
  tft.setCursor(245, 10); tft.print(F("Day ")); tft.println(stats.day);
  tft.setCursor(245, 30); tft.print(F("Re ")); tft.print(stats.current_R,1);
  tft.setCursor(245, 50); tft.print(F("CFR ")); tft.print(stats.case_fatality_rate,1); tft.print(F("%"));
  tft.setCursor(245, 70); tft.print(F("H ")); tft.print(stats.healthcare_capacity); tft.print(F("%"));
  
  byte dominant_strain = 0;
  byte max_count = 0;
  for(byte s=0; s<8; s++){
    if(stats.strain_counts[s] > max_count){
      max_count = stats.strain_counts[s];
      dominant_strain = s;
    }
  }
  tft.setCursor(245, 90); 
  tft.println(strain_cache[dominant_strain].name);
  
  byte active_strains = 0;
  for(byte s=0; s<8; s++) if(stats.strain_counts[s]>0) active_strains++;
  if(active_strains > 1){
    tft.setCursor(245, 110);
    tft.print(active_strains); tft.print(F(" strains"));
  }
  
  // Show CVI
  tft.setCursor(245, 130);
  tft.print(F("CVI ")); tft.print(countryProfile.CVI, 2);
}

void logStatistics(){
  Serial.print(stats.day); Serial.print(',');
  Serial.print(stats.susceptible*populationScale); Serial.print(',');
  Serial.print(stats.exposed*populationScale); Serial.print(',');
  Serial.print(stats.infected*populationScale); Serial.print(',');
  Serial.print(stats.quarantined*populationScale); Serial.print(',');
  Serial.print(stats.recovered*populationScale); Serial.print(',');
  Serial.print(stats.dead*populationScale); Serial.print(',');
  Serial.print(stats.vaccinated*populationScale); Serial.print(',');
  Serial.print(stats.current_R,2); Serial.print(',');
  Serial.print(stats.case_fatality_rate,2); Serial.print(',');
  Serial.print(stats.healthcare_capacity); Serial.print(',');
  
  byte dominant = 0;
  byte max_cnt = 0;
  for(byte s=0; s<8; s++){
    if(stats.strain_counts[s] > max_cnt){
      max_cnt = stats.strain_counts[s];
      dominant = s;
    }
  }
  Serial.print(strain_cache[dominant].name); Serial.print(',');
  
  switch(simState){
    case RUNNING: Serial.println(F("RUNNING")); break;
    case PAUSED_ELIMINATION: Serial.println(F("ELIMINATED")); break;
    case PAUSED_COLLAPSE: Serial.println(F("COLLAPSED")); break;
    case ENDED_ENDEMIC: Serial.println(F("ENDEMIC")); break;
  }
}

//------------------------------------------------------------------------------
// COMMAND PARSING
//------------------------------------------------------------------------------
void parseExactCommand(String line){
  int eq = line.indexOf('=');
  String cmd = line.substring(0,eq); cmd.trim(); cmd.toUpperCase();
  String val = line.substring(eq+1); val.trim();
  
  if(cmd=="I"){
    float pct = val.toFloat();
    transmissionModifier = 1.0 + pct/100.0;
    Serial.print(F("Transmission: +")); Serial.print(pct,1); Serial.println(F("%"));
  }
  else if(cmd=="K"){
    float pct = val.toFloat();
    transmissionModifier = 1.0 - pct/100.0;
    if(transmissionModifier<0.1) transmissionModifier=0.1;
    Serial.print(F("Transmission: -")); Serial.print(pct,1); Serial.println(F("%"));
  }
  else if(cmd=="O"){
    float pct = val.toFloat();
    mortalityModifier = 1.0 + pct/100.0;
    Serial.print(F("Mortality: +")); Serial.print(pct,1); Serial.println(F("%"));
  }
  else if(cmd=="P"){
    float pct = val.toFloat();
    mortalityModifier = 1.0 - pct/100.0;
    if(mortalityModifier<0.1) mortalityModifier=0.1;
    Serial.print(F("Mortality: -")); Serial.print(pct,1); Serial.println(F("%"));
  }
  else if(cmd=="U"){
    if(val.indexOf('%') >= 0){
      float prob = val.substring(0, val.indexOf('%')).toFloat();
      interventions.daily_import_probability = prob / 100.0;
      interventions.import_probability_mode = true;
      Serial.print(F("Daily import probability: ")); Serial.print(prob,2); Serial.println(F("%"));
    } else {
      int count = val.toInt();
      if(count>0 && count<=100){
        importMultiple((byte)count);
        interventions.import_probability_mode = false;
        interventions.daily_import_probability = 0;
        Serial.print(F("Imported ")); Serial.print(count); Serial.println(F(" cases"));
      }
    }
  }
  else if(cmd=="V"){
    int v=val.toInt();
    if(v>=0&&v<=100){
      interventions.vaccine_rate=v;
      interventions.auto_vaccination=false;
      Serial.print(F("Vaccine: ")); Serial.print(v); Serial.println(F("% (manual)"));
    }
  }
  else if(cmd=="T"){
    int t=val.toInt();
    if(t>=0&&t<=25){
      interventions.test_rate=t;
      interventions.auto_testing_scaleup=false;
      Serial.print(F("Testing: ")); Serial.print(t); Serial.println(F("% (manual)"));
    }
  }
  else if(cmd=="L"){
    val.toLowerCase();
    interventions.lockdown=(val=="on");
    Serial.print(F("Lockdown: ")); Serial.println(interventions.lockdown?F("ON"):F("OFF"));
  }
  else if(cmd=="M"){
    val.toLowerCase();
    interventions.mask_mandate=(val=="on");
    Serial.print(F("Masks: ")); Serial.println(interventions.mask_mandate?F("ON"):F("OFF"));
  }
  else if(cmd=="AUTO"){
    val.toLowerCase();
    bool st=(val=="on");
    interventions.auto_behavioral_response=st;
    interventions.auto_testing_scaleup=st;
    interventions.auto_vaccination=st;
    interventions.auto_healthcare_crisis=st;
    interventions.auto_seasonal_forcing=st;
    interventions.auto_compliance_fatigue=st;
    interventions.auto_variant_emergence=st;
    interventions.auto_importation=st;
    Serial.print(F("Autonomous systems: ")); Serial.println(st?F("ON"):F("OFF"));
  }
  else if(cmd=="DEMO"){
    val.toLowerCase();
    DEMO_MODE=(val=="on");
    Serial.print(F("Demo mode: ")); Serial.println(DEMO_MODE?F("ON"):F("OFF"));
  }
  else if(cmd=="SCENARIO"){
    int s=val.toInt();
    if(s>=1&&s<=5) loadScenario(s);
  }
  else {
    Serial.println(F("Unknown command"));
  }
}

void setScaleByCountryName(String name){
  name.toLowerCase();
  
  for(byte i=0; i<NUM_COUNTRIES; i++){
    char country_name[16];
    memcpy_P(country_name, countries[i].name, 16);
    String cname = String(country_name);
    cname.toLowerCase();
    
    if(name == cname){
      activeCountry = i;
      loadCountryProfile(i);
      return;
    }
  }
  
  Serial.println(F("Unknown country. Available:"));
  Serial.println(F("  New Zealand, Germany, Australia, Canada, USA,"));
  Serial.println(F("  UK, Spain, China, Russia, Iran, Global"));
}

void loadScenario(byte s){
  switch(s){
    case 1:
      Serial.println(F("Scenario: Do Nothing"));
      interventions.lockdown=false;
      interventions.mask_mandate=false;
      interventions.vaccine_rate=0;
      interventions.test_rate=0;
      interventions.auto_vaccination=false;
      interventions.auto_testing_scaleup=false;
      break;
    case 2:
      Serial.println(F("Scenario: Early Lockdown"));
      interventions.lockdown=true;
      interventions.mask_mandate=true;
      interventions.test_rate=20;
      interventions.compliance=0.90;
      break;
    case 3:
      Serial.println(F("Scenario: Mass Vaccination"));
      interventions.vaccine_rate=3;
      interventions.prioritize_elderly=true;
      interventions.auto_vaccination=true;
      break;
    case 4:
      Serial.println(F("Scenario: Intermittent"));
      interventions.auto_behavioral_response=true;
      interventions.compliance=0.70;
      break;
    case 5:
      Serial.println(F("Scenario: Elimination"));
      interventions.lockdown=true;
      interventions.test_rate=25;
      interventions.compliance=0.95;
      interventions.auto_importation=false;
      break;
  }
}

//------------------------------------------------------------------------------
// RESET
//------------------------------------------------------------------------------
void resetSimulation(){
  for(int i=0;i<TOTAL_POP;i++){
    byte age; byte r=random(100);
    if(r<25) age=0; else if(r<85) age=1; else age=2;
    
    population[i].state=SUSCEPTIBLE;
    population[i].age_group=age;
    population[i].pos_row=i/GRID_DIM;
    population[i].pos_col=i%GRID_DIM;
    population[i].daysExposed=0; 
    population[i].latent_duration=0;
    population[i].daysInfected=0; 
    population[i].infectious_duration=0;
    population[i].immunity_level=0; 
    population[i].days_since_recovery=0;
    population[i].infecting_strain=currentStrain;
  }

  byte startRow = GRID_DIM/2 - 1;
  byte startCol = GRID_DIM/2 - 1;
  
  for(byte dr=0; dr<2; dr++){
    for(byte dc=0; dc<2; dc++){
      int idx = (startRow+dr)*GRID_DIM + (startCol+dc);
      population[idx].state = EXPOSED;
      population[idx].daysExposed = random(1,3);
      population[idx].latent_duration = 5;
      population[idx].infecting_strain = currentStrain;
    }
  }

  memset(&stats,0,sizeof(stats));
  stats.susceptible = TOTAL_POP - 4;
  stats.exposed = 4;
  stats.total_exposed = 4;
  stats.current_R = calculateEffectiveR();
  
  for(byte s=0; s<8; s++){
    stats.strain_counts[s] = 0;
    stats.strain_frequency[s] = 0;
  }
  stats.strain_counts[currentStrain] = 4;
  stats.strain_frequency[currentStrain] = 1.0;

  // Apply country-specific CVI modifiers
  transmissionModifier = countryProfile.transmission_modifier;
  mortalityModifier = countryProfile.mortality_modifier;
  
  interventions.lockdown = false;
  interventions.school_closure = false;
  interventions.event_ban = false;
  interventions.mask_mandate = false;
  interventions.test_rate = 0; 
  interventions.vaccine_rate = 0;
  interventions.compliance = 0.75;
  interventions.behavioral_reduction = 0.0;
  interventions.healthcare_strain = 1.0; 
  interventions.seasonal_modifier = 1.0;
  interventions.compliance_decay = 1.0; 
  interventions.days_since_last_intervention = 0;
  interventions.daily_import_probability = 0;
  interventions.import_probability_mode = false;
  
  interventions.auto_behavioral_response = true; 
  interventions.auto_testing_scaleup = true;
  interventions.auto_vaccination = true; 
  interventions.auto_healthcare_crisis = true;
  interventions.auto_seasonal_forcing = true; 
  interventions.auto_compliance_fatigue = true;
  interventions.auto_variant_emergence = true; 
  interventions.auto_importation = true;
  
  announced_variants = 0;  // FIX #2: Reset on simulation restart
  simState = RUNNING;

  lcd.clear(); 
  lcd.print(F("RESET")); 
  lcd.setCursor(0,1); 
  lcd.print(disease.name);
  delay(800);
}

//------------------------------------------------------------------------------
// EVENT DETECTION
//------------------------------------------------------------------------------
void checkEpidemicEvents(){
  static bool epi=false, peak=false, vax_started=false, half_vax=false, lockdown_active=false;
  
  float prev = ((float)(stats.infected+stats.exposed)/TOTAL_POP)*100;
  if(!epi && prev>1.0){
    Serial.println(F("\n*** EPIDEMIC DECLARED (>1% prevalence) ***\n"));
    playTone(1500,300); 
    epi=true;
  }
  
  static byte declining=0; 
  static byte last_inf=0;
  if(stats.infected<last_inf) declining++; else declining=0;
  last_inf=stats.infected;
  if(!peak && declining>=7 && stats.total_infected>10){
    Serial.println(F("\n*** EPIDEMIC PEAK REACHED ***\n"));
    playTone(800,200); 
    peak=true;
  }
  
  if(!vax_started && interventions.vaccine_rate>0.5){
    Serial.println(F("\n*** VACCINATION CAMPAIGN LAUNCHED ***\n"));
    playTone(1800,100); delay(50); playTone(2000,100); delay(50); playTone(2200,200);
    vax_started=true;
  }
  
  if(!half_vax && stats.vaccinated > TOTAL_POP/2){
    Serial.println(F("\n*** 50% VACCINATED ***\n"));
    playTone(1500,300);
    half_vax=true;
  }
  
  if(!lockdown_active && interventions.lockdown){
    Serial.println(F("\n*** LOCKDOWN IN EFFECT ***\n"));
    playTone(600,400);
    lockdown_active=true;
  }
  if(lockdown_active && !interventions.lockdown){
    Serial.println(F("\n*** LOCKDOWN LIFTED ***\n"));
    playTone(1000,200);
    lockdown_active=false;
  }
}

//------------------------------------------------------------------------------
// UTILITIES
//------------------------------------------------------------------------------
void loadAllStrains(){
  for(byte i=0; i<8; i++){
    memcpy_P(&strain_cache[i], &strains[i], sizeof(DiseaseStrain));
  }
}

void loadCountryProfile(byte idx){
  memcpy_P(&countryProfile, &countries[idx], sizeof(CountryProfile));
  populationScale = countryProfile.population / TOTAL_POP;
  
  Serial.print(F("\n--- ")); Serial.print(countryProfile.name); Serial.println(F(" LOADED ---"));
  Serial.print(F("Population: ")); Serial.println(countryProfile.population);
  Serial.print(F("CVI: ")); Serial.print(countryProfile.CVI, 2); 
  Serial.println(F(" (1.00 = global average)"));
  Serial.print(F("Transmission modifier: ")); Serial.print(countryProfile.transmission_modifier, 2);
  Serial.print(F(" (")); Serial.print((countryProfile.transmission_modifier-1)*100, 1); Serial.println(F("%)"));
  Serial.print(F("Mortality modifier: ")); Serial.print(countryProfile.mortality_modifier, 2);
  Serial.print(F(" (")); Serial.print((countryProfile.mortality_modifier-1)*100, 1); Serial.println(F("%)\n"));
}

float getContactRate(byte age1, byte age2){
  return pgm_read_float(&contact_matrix[age1][age2]);
}

void playTone(unsigned int freq, unsigned int dur){
  tone(BUZZER_PIN, freq, dur);
  delay(dur);
  noTone(BUZZER_PIN);
}

//------------------------------------------------------------------------------
// RANDOM NUMBER GENERATORS
//------------------------------------------------------------------------------
float randomGamma(float shape, float scale){
  if(shape<1.0){
    float u,v,x,y;
    while(1){
      u=random(10000)/10000.0; v=random(10000)/10000.0;
      if(u<=1.0-shape){ x=pow(u,1.0/shape); y=v*pow(x,shape-1.0); }
      else{ x=1.0-log((u-(1.0-shape))/shape); y=v*exp(-x); }
      if(y<=pow(x,shape-1.0)*exp(-x)) return x*scale;
    }
  }else{
    float d=shape-1.0/3.0, c=1.0/sqrt(9.0*d);
    while(1){
      float x,v;
      do{
        float u1=random(10000)/10000.0; if(u1<0.0001)u1=0.0001;
        float u2=random(10000)/10000.0;
        x=sqrt(-2.0*log(u1))*cos(2.0*PI*u2);
        v=1.0+c*x;
      }while(v<=0.0);
      v=v*v*v;
      float u=random(10000)/10000.0;
      if(u<1.0-0.0331*x*x*x*x) return d*v*scale;
      if(log(u)<0.5*x*x+d*(1.0-v+log(v))) return d*v*scale;
    }
  }
}

byte sampleNegativeBinomial(float mean, float k){
  float lambda = randomGamma(k, mean/k);
  return randomPoisson(lambda);
}

byte randomPoisson(float lambda){
  if(lambda<=0) return 0;
  if(lambda>30) return 30;
  float L=exp(-lambda); byte k=0; float p=1.0;
  do{ k++; p*=random(1000)/1000.0; }while(p>L && k<30);
  return k-1;
}