import numpy as np
from sklearn.linear_model import LinearRegression

# Historical data for each state (years and corresponding elephant population)
data = {
    'Arunachal Pradesh': np.array([2102, 1800, 1607, 1690, 890, 1614]),
    'Assam': np.array([5524, 5312, 5246, 5281, 5620, 5719]),
    'Meghalaya': np.array([2872, 1840, 1868, 1811, 1811, 1754]),
    'Nagaland': np.array([178, 158, 145, 152, 212, 446]),
    'Mizoram': np.array([15, 22, 33, 12, 10, 7]),
    'Manipur': np.array([50, 30, 12, 11, 12, 9]),
    'Tripura': np.array([100, 70, 40, 59, 59, 102]),
    'West Bengal (North)': np.array([186, 250, 292, 300, 647, 488]),
    'West Bengal (South)': np.array([14, 26, 36, 25, 109, 194]),
    'Jharkhand': np.array([550, 618, 772, 624, 688, 679]),
    'Odisha': np.array([1750, 1800, 1841, 1862, 1930, 1976]),
    'Chhattisgarh': np.array([124, 123, 122, 122, 247, 247]),
    'Bihar': np.array([69, 62, 55, 48, 41, 25]),
    'Madhya Pradesh': np.array([9, 8, 7, 7, 7, 7]),
    'Uttarakhand': np.array([828, 1130, 1582, 1346, 1346, 1839]),
    'Uttar Pradesh': np.array([47, 70, 85, 380, 291, 232]),
    'Haryana': np.array([9, 8, 7, 7, 7, 7]),
    'Himachal Pradesh': np.array([9, 8, 7, 7, 7, 7]),
    'Tamil Nadu': np.array([2307, 2971, 3052, 3867, 4015, 2761]),
    'Karnataka': np.array([5500, 6088, 5838, 4035, 6068, 6049]),
    'Kerala': np.array([3500, 3600, 3850, 6068, 6185, 5706]),
    'Andhra Pradesh': np.array([46, 57, 74, 28, 41, 65]),
    'Maharashtra': np.array([9, 8, 7, 7, 4, 6]),
    'A&N Nicobar Islands': np.array([9, 8, 7, 7, 7, 25])
}

# Create a linear regression model for each state
models = {}
for state, state_data in data.items():
    years = np.array([1993, 1997, 2002, 2007, 2012, 2017])
    population = state_data

    model = LinearRegression()
    model.fit(years.reshape(-1, 1), population)
    models[state] = model

# Predict the elephant population for the years 2018 to 2022 for each state
future_years = np.array([2018, 2019, 2020, 2021, 2022])
predictions = {}

for state, model in models.items():
    predicted_population = model.predict(future_years.reshape(-1, 1))
    predictions[state] = predicted_population

# Print the predicted population for each state and year
for state, predicted_population in predictions.items():
    print(f"State: {state}")
    for year, population in zip(future_years, predicted_population):
        print(f"Year: {year}, Predicted Population: {int(population)}")
    print()