import numpy as np
from scipy.stats import linregress
from netCDF4 import Dataset

def calculate_and_save_slope(input_file_path, output_file_path):
    # Load the existing netCDF file
    dataset = Dataset(input_file_path, 'r')  # Open the file in read mode

    # Extract the variable values
    time = dataset.variables['time'][:]  # Read the time variable
    ndvi = dataset.variables['Temperature_Air_2m_Mean_24h'][:]  # Extract the NDVI data stored as 'var167'

    # Create a scaled time variable (from 1 to 41)
    scaled_time = np.arange(1, len(time) + 1)

    # Get the dimensions
    num_time_steps, num_lat, num_lon = ndvi.shape  # Extract time, latitude, and longitude dimensions of NDVI data

    # Create arrays to store the slope, intercept, and other information
    slope_values = np.zeros((num_lat, num_lon))
    intercept_values = np.zeros((num_lat, num_lon))
    r_value_values = np.zeros((num_lat, num_lon))
    p_value_values = np.zeros((num_lat, num_lon))
    std_err_values = np.zeros((num_lat, num_lon))

    # Perform linear regression for each cell
    for i in range(num_lat):
        for j in range(num_lon):
            # Check if any NDVI value is nan, if yes, set the slope value to nan
            if np.isnan(ndvi[:, i, j]).any():
                slope_values[i, j] = np.nan
                intercept_values[i, j] = np.nan
                r_value_values[i, j] = np.nan
                p_value_values[i, j] = np.nan
                std_err_values[i, j] = np.nan
            else:
                slope, intercept, r_value, p_value, std_err = linregress(scaled_time, ndvi[:, i, j])
                slope_values[i, j] = slope
                intercept_values[i, j] = intercept
                r_value_values[i, j] = r_value
                p_value_values[i, j] = p_value
                std_err_values[i, j] = std_err

    # Close the input file
    dataset.close()

    # Create a new netCDF file with a different name
    output_dataset = Dataset(output_file_path, 'w', format='NETCDF4')

    # Define dimensions in the new file
    output_dataset.createDimension('lat', num_lat)
    output_dataset.createDimension('lon', num_lon)

    # Create new variables for slope, intercept, r_value, p_value, and std_err in the new dataset
    slope_var = output_dataset.createVariable('slope', 'f4', dimensions=('lat', 'lon'))
    slope_var[:, :] = slope_values

    intercept_var = output_dataset.createVariable('intercept', 'f4', dimensions=('lat', 'lon'))
    intercept_var[:, :] = intercept_values

    r_value_var = output_dataset.createVariable('r_value', 'f4', dimensions=('lat', 'lon'))
    r_value_var[:, :] = r_value_values

    p_value_var = output_dataset.createVariable('p_value', 'f4', dimensions=('lat', 'lon'))
    p_value_var[:, :] = p_value_values

    std_err_var = output_dataset.createVariable('std_err', 'f4', dimensions=('lat', 'lon'))
    std_err_var[:, :] = std_err_values

    # Close the new file
    output_dataset.close()

    print(f"Regression information saved to the new netCDF file: {output_file_path}")

input_path = '/temperature_anomaly.nc'
output_path = '/temperature_liner_regression.nc'
calculate_and_save_slope(input_path, output_path)