use bevy::color::palettes::css;
use bevy::prelude::{ops::sin,ops::powf,*};
use bevy::render::{
    render_asset::RenderAssetUsages,
    render_resource::{Extent3d, TextureDimension, TextureFormat},
};

const IMAGE_WIDTH: usize = 200;
const IMAGE_HEIGHT: usize = 100;
const SCREEN_WIDTH: f32 = 2000.; 
const SCREEN_HEIGHT: f32 = 1000.;
const PML_THICKNESS: usize = 10;
const MAX_CONDUCTIVITY: f32 = 1.0;
const S_C: f32 = 0.01;

fn main() {
    App::new()
        .add_plugins(DefaultPlugins.set(WindowPlugin {
                primary_window: Some(Window {
                    resolution: (SCREEN_WIDTH, SCREEN_HEIGHT).into(),
                    ..default()
                }),
                ..default()
        }))
        .insert_resource(Time::<Fixed>::from_hz(24.0))
        .add_systems(Startup, setup)
        .add_systems(FixedUpdate, update_fields)
        .add_systems(FixedUpdate, radiate_point_sources)
        .run();
}

#[derive(Resource)]
struct Grid{
    handle: Handle<Image>,
    e_z: Vec<Vec<f32>>,
    e_zx: Vec<Vec<f32>>,
    e_zy: Vec<Vec<f32>>,
    b_x: Vec<Vec<f32>>,
    b_y: Vec<Vec<f32>>,
    eps: Vec<Vec<f32>>,
    mu: Vec<Vec<f32>>,
}


#[derive(Component)]
struct PointSource{
    x: usize,
    y: usize,
    omega: f32,
    amplitude: f32
}

fn setup(mut commands: Commands, mut images: ResMut<Assets<Image>>) {
    commands.spawn(Camera2d);

    let image = Image::new_fill(
        Extent3d {
            width: IMAGE_WIDTH as u32,
            height: IMAGE_HEIGHT as u32,
            depth_or_array_layers: 1,
        },
        TextureDimension::D2,
        &(css::WHITE.to_u8_array()),
        TextureFormat::Rgba8UnormSrgb,
        RenderAssetUsages::MAIN_WORLD | RenderAssetUsages::RENDER_WORLD,
    );

    let h = images.add(image);

    commands.spawn(Sprite {
        image: h.clone(),
        custom_size: Some(Vec2::new(SCREEN_WIDTH, SCREEN_HEIGHT)),
        ..default()
    });
    commands.insert_resource(Grid{
        handle: h,
        e_z: vec![vec![0.0; IMAGE_HEIGHT]; IMAGE_WIDTH],
        e_zx: vec![vec![0.0; IMAGE_HEIGHT]; IMAGE_WIDTH],
        e_zy: vec![vec![0.0; IMAGE_HEIGHT]; IMAGE_WIDTH],
        b_x: vec![vec![0.0; IMAGE_HEIGHT]; IMAGE_WIDTH],
        b_y: vec![vec![0.0; IMAGE_HEIGHT]; IMAGE_WIDTH],
        eps: vec![vec![1.0; IMAGE_HEIGHT]; IMAGE_WIDTH],
        mu: vec![vec![1.0; IMAGE_HEIGHT]; IMAGE_WIDTH],
    });
    commands.spawn(PointSource{amplitude: 3.0, omega: 2.0, x: 100, y: 50});
}

fn update_fields(mut grid: ResMut<Grid>, mut images: ResMut<Assets<Image>>) {
    let image = images.get_mut(&grid.handle).expect("Image not found");
    for x in 1..(IMAGE_WIDTH-1) {
        for y in 1..(IMAGE_HEIGHT-1) {
            let eps = grid.eps[x][y]; 
            let mu = grid.mu[x][y];

            if x < PML_THICKNESS || IMAGE_WIDTH - x < PML_THICKNESS || y < PML_THICKNESS || IMAGE_HEIGHT - y < PML_THICKNESS {
                let (sigma_x,sigma_y) = get_conductivity(x,y);
                let ex = (mu/S_C-sigma_x/2.0)/(mu/S_C+sigma_x/2.0);
                let fx = 1.0/(mu/S_C+sigma_x/2.0);
                let ey = (mu/S_C-sigma_y/2.0)/(mu/S_C+sigma_y/2.0);
                let fy = 1.0/(mu/S_C+sigma_y/2.0);
                grid.b_x[x][y] += ey*-1.0*(grid.e_z[x][y+1]-grid.e_z[x][y]) - fy*grid.b_x[x][y];
                grid.b_y[x][y] += ex*(grid.e_z[x+1][y]-grid.e_z[x][y]) - fx*grid.b_y[x][y];
                let cx = (eps/S_C-sigma_x/2.0)/(eps/S_C+sigma_x/2.0);
                let dx = 1.0/(eps/S_C+sigma_x/2.0);
                let cy = (eps/S_C-sigma_y/2.0)/(eps/S_C+sigma_y/2.0);
                let dy = 1.0/(eps/S_C+sigma_y/2.0);
                grid.e_zx[x][y] += cx*(grid.b_x[x][y-1] - grid.b_x[x][y]) - dx*grid.e_zx[x][y];
                grid.e_zy[x][y] += cy*-1.0*(grid.b_y[x-1][y] - grid.b_y[x][y]) - dy*grid.e_zy[x][y];
                grid.e_z[x][y] = grid.e_zx[x][y] + grid.e_zy[x][y];
            } else {
                grid.b_x[x][y] += (S_C/mu) * (-1.0*(grid.e_z[x][y+1]-grid.e_z[x][y]));
                grid.b_y[x][y] += (S_C/mu) * (grid.e_z[x+1][y]-grid.e_z[x][y]);
                grid.e_z[x][y] += (S_C/eps) * ((grid.b_x[x][y-1] - grid.b_x[x][y]) - (grid.b_y[x-1][y] - grid.b_y[x][y]));
            }
            if grid.e_z[x][y] > 0.0 {
                let col = Color::srgba(grid.e_z[x][y] * grid.e_z[x][y], 0.0, 0.0, 1.0);
                let _ = image.set_color_at(x as u32, y as u32, col);
            } else {
                let col = Color::srgba(0.0, grid.e_z[x][y] * grid.e_z[x][y], 0.0, 1.0);
                let _ = image.set_color_at(x as u32, y as u32, col);
            }
        }
    }
}

fn radiate_point_sources(mut grid: ResMut<Grid>, sources: Query<&PointSource>, time: Res<Time<Fixed>>) {
    for source in sources {
        grid.e_z[source.x][source.y] = source.amplitude*sin(source.omega*time.elapsed().as_secs_f32())
    }
}

fn get_conductivity(x: usize,y: usize) -> (f32,f32) {
    let mut sigma_x: f32 = 0.0;
    let mut sigma_y: f32 = 0.0;
    let x_f = x as f32;
    let y_f = y as f32;
    let d = PML_THICKNESS as f32;
    if x < PML_THICKNESS {
        sigma_x = MAX_CONDUCTIVITY*powf(x_f/d, 1.0);
    }
    if IMAGE_WIDTH - x < PML_THICKNESS {
        sigma_x = MAX_CONDUCTIVITY*powf((IMAGE_WIDTH as f32 - x_f)/d, 3.0);
    }
    if y < PML_THICKNESS {
        sigma_y = MAX_CONDUCTIVITY*powf(y_f/d, 2.0);
    }
    if IMAGE_HEIGHT- y < PML_THICKNESS {
        sigma_y = MAX_CONDUCTIVITY*powf((IMAGE_HEIGHT as f32 - y_f)/d, 3.0);
    }
    return (sigma_x,sigma_y);
}