import '../models/models.dart';

/// Simplified mortality groups with approximate life expectancy at birth.
/// Data approximated from WHO 2024 estimates.
enum MortalityGroup {
  groupA, // High LE countries (~83-85)
  groupB, // Upper-middle LE (~79-81)
  groupC, // Middle LE (~72-76)
  groupD, // Lower LE (~65-70)
}

/// Maps countries to their mortality group.
MortalityGroup getCountryGroup(String country) {
  return _countryToGroup[country] ?? MortalityGroup.groupB;
}

/// Get baseline life expectancy at birth for a given country group and sex.
double getLifeExpectancyAtBirth(MortalityGroup group, Sex sex) {
  final base = _groupBaseLE[group]!;
  // Women live ~4-5 years longer on average
  return sex == Sex.female ? base + 4.5 : base;
}

/// Get remaining life expectancy at current age.
/// Simplified model: as you age, remaining LE decreases but survivors
/// tend to live longer than birth LE suggests.
double getRemainingLifeExpectancy(int currentAge, Sex sex, String country) {
  final group = getCountryGroup(country);
  final leAtBirth = getLifeExpectancyAtBirth(group, sex);

  if (currentAge >= leAtBirth) {
    // Past average LE - use simplified survival model
    // Each year survived past LE adds ~0.5-0.8 expected years
    return 5.0 + (leAtBirth - currentAge) * 0.1;
  }

  // Simplified remaining LE calculation
  // People who survive to age X have higher LE than birth cohort suggests
  final survivorBonus = currentAge * 0.15; // ~0.15 years bonus per year survived
  final rawRemaining = leAtBirth - currentAge;

  return rawRemaining + survivorBonus.clamp(0, 5);
}

/// Base life expectancy by mortality group (male baseline).
const _groupBaseLE = {
  MortalityGroup.groupA: 81.0,
  MortalityGroup.groupB: 77.0,
  MortalityGroup.groupC: 72.0,
  MortalityGroup.groupD: 65.0,
};

/// Country to mortality group mapping.
const _countryToGroup = {
  // Group A - High LE (83-85)
  'Japan': MortalityGroup.groupA,
  'Switzerland': MortalityGroup.groupA,
  'Singapore': MortalityGroup.groupA,
  'Spain': MortalityGroup.groupA,
  'Italy': MortalityGroup.groupA,
  'Australia': MortalityGroup.groupA,
  'Iceland': MortalityGroup.groupA,
  'Israel': MortalityGroup.groupA,
  'Sweden': MortalityGroup.groupA,
  'France': MortalityGroup.groupA,
  'South Korea': MortalityGroup.groupA,
  'Norway': MortalityGroup.groupA,

  // Group B - Upper-middle LE (79-81)
  'United States': MortalityGroup.groupB,
  'United Kingdom': MortalityGroup.groupB,
  'Germany': MortalityGroup.groupB,
  'Canada': MortalityGroup.groupB,
  'Netherlands': MortalityGroup.groupB,
  'Belgium': MortalityGroup.groupB,
  'Austria': MortalityGroup.groupB,
  'Finland': MortalityGroup.groupB,
  'Ireland': MortalityGroup.groupB,
  'New Zealand': MortalityGroup.groupB,
  'Denmark': MortalityGroup.groupB,
  'Portugal': MortalityGroup.groupB,
  'Czech Republic': MortalityGroup.groupB,
  'Poland': MortalityGroup.groupB,
  'Chile': MortalityGroup.groupB,
  'Costa Rica': MortalityGroup.groupB,
  'Cuba': MortalityGroup.groupB,
  'United Arab Emirates': MortalityGroup.groupB,
  'Qatar': MortalityGroup.groupB,
  'Taiwan': MortalityGroup.groupB,

  // Group C - Middle LE (72-76)
  'China': MortalityGroup.groupC,
  'Brazil': MortalityGroup.groupC,
  'Mexico': MortalityGroup.groupC,
  'Russia': MortalityGroup.groupC,
  'Turkey': MortalityGroup.groupC,
  'Argentina': MortalityGroup.groupC,
  'Colombia': MortalityGroup.groupC,
  'Thailand': MortalityGroup.groupC,
  'Vietnam': MortalityGroup.groupC,
  'Malaysia': MortalityGroup.groupC,
  'Iran': MortalityGroup.groupC,
  'Saudi Arabia': MortalityGroup.groupC,
  'Egypt': MortalityGroup.groupC,
  'Ukraine': MortalityGroup.groupC,
  'Romania': MortalityGroup.groupC,
  'Hungary': MortalityGroup.groupC,
  'Peru': MortalityGroup.groupC,
  'Philippines': MortalityGroup.groupC,

  // Group D - Lower LE (65-70)
  'India': MortalityGroup.groupD,
  'Indonesia': MortalityGroup.groupD,
  'South Africa': MortalityGroup.groupD,
  'Pakistan': MortalityGroup.groupD,
  'Bangladesh': MortalityGroup.groupD,
  'Nigeria': MortalityGroup.groupD,
  'Kenya': MortalityGroup.groupD,
  'Ghana': MortalityGroup.groupD,
  'Ethiopia': MortalityGroup.groupD,
  'Myanmar': MortalityGroup.groupD,
  'Nepal': MortalityGroup.groupD,
  'Cambodia': MortalityGroup.groupD,
};

/// Get list of all supported countries, sorted alphabetically.
List<String> getSupportedCountries() {
  final countries = _countryToGroup.keys.toList();
  countries.sort();
  return countries;
}