use bevy::color::palettes::css;
use bevy::prelude::*;
use bevy::prelude::ops::{sin,powf,exp};
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 EPS_0: f32 = 1.0;
const DT: f32 = 0.1;
const DR: f32 = 1.0;

const M: f32 = 3.0;
const PML_THICKNESS: usize = 20;
const SIM_WIDTH: usize = IMAGE_WIDTH+2*PML_THICKNESS;
const SIM_HEIGHT: usize = IMAGE_HEIGHT+2*PML_THICKNESS;
const MAX_CONDUCTIVITY: f32 = -1.0*(M+1.0)*(-9.2)/2.0/(PML_THICKNESS as f32);
const KAPPA_MAX: f32 = 2.0;
const A_MAX: f32 = 1.0;

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)
        .add_systems(Update, mouse_controls)
        .add_systems(Update, keyboard_controls)
        .run();
}

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


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

#[derive(Component)]
enum UserState{
    SelectPointSource,
    SelectWall,
    SelectFreeSpace,
    SelectNothing
}

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; SIM_HEIGHT]; SIM_WIDTH],
        psi_ex: vec![vec![0.0; SIM_HEIGHT]; SIM_WIDTH],
        psi_ey: vec![vec![0.0; SIM_HEIGHT]; SIM_WIDTH],
        b_x: vec![vec![0.0; SIM_HEIGHT]; SIM_WIDTH],
        b_y: vec![vec![0.0; SIM_HEIGHT]; SIM_WIDTH],
        psi_bx: vec![vec![0.0; SIM_HEIGHT]; SIM_WIDTH],
        psi_by: vec![vec![0.0; SIM_HEIGHT]; SIM_WIDTH],
        eps: vec![vec![1.0; SIM_HEIGHT]; SIM_WIDTH],
        mu: vec![vec![1.0; SIM_HEIGHT]; SIM_WIDTH],
    });
    commands.spawn(UserState::SelectNothing);
}

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..(SIM_WIDTH-1) {
        for y in 1..(SIM_HEIGHT-1) {
            let eps = grid.eps[x][y]; 
            let mu = grid.mu[x][y];

            if x < PML_THICKNESS || SIM_WIDTH < (PML_THICKNESS+x) || y < PML_THICKNESS || SIM_HEIGHT < (PML_THICKNESS+y) {
                let (px,py) = p(x,y);
                let (qx,qy) = q(x,y);
                let (kx,ky) = get_kappa(x, y);
                grid.psi_bx[x][y] = px*grid.psi_bx[x][y] + qx*(grid.e_z[x+1][y] - grid.e_z[x][y])/DR;
                grid.psi_by[x][y] = py*grid.psi_by[x][y] + qy*(grid.e_z[x][y+1] - grid.e_z[x][y])/DR;
                grid.b_x[x][y] += (-1.0*DT/mu) * ((grid.e_z[x][y+1] - grid.e_z[x][y])/DR/ky + grid.psi_by[x][y]);
                grid.b_y[x][y] += (DT/mu) * ((grid.e_z[x+1][y] - grid.e_z[x][y])/DR/kx + grid.psi_bx[x][y]);

                grid.psi_ex[x][y] = px*grid.psi_ex[x][y] + qx*(grid.b_x[x][y-1] - grid.b_x[x][y])/DR;
                grid.psi_ey[x][y] = py*grid.psi_ey[x][y] + qy*(grid.b_y[x-1][y] - grid.b_y[x][y])/DR;
                grid.e_z[x][y] += (DT/eps) * ((grid.b_x[x][y-1] - grid.b_x[x][y])/DR/ky - (grid.b_y[x-1][y] - grid.b_y[x][y])/DR/kx + (grid.psi_ex[x][y] - grid.psi_ey[x][y]));
            } else {
                grid.b_x[x][y] += (DT/DR/mu) * -1.0*(grid.e_z[x][y+1]-grid.e_z[x][y]);
                grid.b_y[x][y] += (DT/DR/mu) * (grid.e_z[x+1][y]-grid.e_z[x][y]);
                grid.e_z[x][y] += (DT/DR/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].powf(2.), 0.0, eps/10., 1.0);
                    let _ = image.set_color_at((x-PML_THICKNESS) as u32, (y-PML_THICKNESS) as u32, col);
                } else {
                    let col = Color::srgba(0.0, grid.e_z[x][y].powf(2.), eps/10., 1.0);
                    let _ = image.set_color_at((x-PML_THICKNESS) as u32, (y-PML_THICKNESS) 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 p(x :usize,y :usize) -> (f32,f32) {
    let (sigma_x,sigma_y) = get_conductivity(x, y);
    let (kappa_x,kappa_y) = get_kappa(x, y);
    let (a_x,a_y) = get_a(x, y);
    
    if sigma_x == 0.0 {
        return (1.0, exp(-(sigma_y/kappa_y+a_y)*DT/EPS_0) );
    } else if sigma_y == 0.0 {
        return (exp(-(sigma_x/kappa_x+a_x)*DT/EPS_0),1.0);
    } else {
        return (exp(-(sigma_x/kappa_x+a_x)*DT/EPS_0), exp(-(sigma_y/kappa_y+a_y)*DT/EPS_0) );
    }
}

fn q(x :usize,y :usize) -> (f32,f32) {
    let (sigma_x,sigma_y) = get_conductivity(x, y);
    let (kappa_x,kappa_y) = get_kappa(x, y);
    let (a_x,a_y) = get_a(x, y);
    let (b_x,b_y) = p(x,y);

    if sigma_x == 0.0 {
        return (0.0,(sigma_y*(b_y-1.))/(sigma_y*kappa_y+kappa_y.powf(2.)*a_y));
    } else if sigma_y == 0.0 {
        return ((sigma_x*(b_x-1.))/(sigma_x*kappa_x+kappa_x.powf(2.)*a_x),0.0);
    } else {
        return ((sigma_x*(b_x-1.))/(sigma_x*kappa_x+kappa_x.powf(2.)*a_x),(sigma_y*(b_y-1.))/(sigma_y*kappa_y+kappa_y.powf(2.)*a_y));
    }
}

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, M);
    }
    if IMAGE_WIDTH < (PML_THICKNESS+x) {
        sigma_x = MAX_CONDUCTIVITY*powf((IMAGE_WIDTH as f32 - x_f)/d, M);
    }
    if y < PML_THICKNESS {
        sigma_y = MAX_CONDUCTIVITY*powf(y_f/d, M);
    }
    if SIM_HEIGHT < (PML_THICKNESS+y) {
        sigma_y = MAX_CONDUCTIVITY*powf((SIM_HEIGHT as f32 - y_f)/d, M);
    }
    return (sigma_x,sigma_y);
}

fn get_kappa(x: usize,y: usize) -> (f32,f32) {
    let mut kappa_x: f32 = 1.0;
    let mut kappa_y: f32 = 1.0;
    let x_f = x as f32;
    let y_f = y as f32;
    let d = PML_THICKNESS as f32;
    if x < PML_THICKNESS {
        kappa_x = 1.+(KAPPA_MAX-1.)*powf(x_f/d, M);
    }
    if SIM_WIDTH < (PML_THICKNESS+x) {
        kappa_x = 1.+(KAPPA_MAX-1.)*powf((SIM_WIDTH as f32 - x_f)/d, M);
    }
    if y < PML_THICKNESS {
        kappa_y = 1.+(KAPPA_MAX-1.)*powf(y_f/d, M);
    }
    if SIM_HEIGHT < (PML_THICKNESS+y) {
        kappa_y = 1.+(KAPPA_MAX-1.)*powf((SIM_HEIGHT as f32 - y_f)/d, M);
    }
    return (kappa_x,kappa_y);
}

fn get_a(x: usize,y: usize) -> (f32,f32) {
    let mut a_x: f32 = 0.0;
    let mut a_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 {
        a_x = A_MAX*powf((d-x_f)/d, M);
    }
    if SIM_WIDTH < (PML_THICKNESS+x) {
        a_x = A_MAX*powf((d-(SIM_WIDTH as f32 - x_f))/d, M);
    }
    if y < PML_THICKNESS {
        a_y = A_MAX*powf((d-y_f)/d, M);
    }
    if SIM_HEIGHT < (PML_THICKNESS+y) {
        a_y = A_MAX*powf((d-(SIM_HEIGHT as f32 - y_f))/d, M);
    }
    return (a_x,a_y);
}

fn mouse_controls(window: Query<&Window>, user_state: Query<&UserState>, buttons: Res<ButtonInput<MouseButton>>, mut grid: ResMut<Grid>, mut commands: Commands) {
    let position = window.single().unwrap().cursor_position().unwrap_or(vec2(-1., -1.));

    if buttons.just_pressed(MouseButton::Left) {
        for state in user_state {
            match state {
                UserState::SelectPointSource => {
                    if position != vec2(-1., -1.) {
                        let x = (position.x/10.0)as usize + PML_THICKNESS;
                        let y = (position.y/10.0)as usize + PML_THICKNESS;
                        commands.spawn(PointSource{amplitude: 3.0, omega: 1.0, x, y});
                    }
                }
                UserState::SelectWall => {
                    if position != vec2(-1., -1.) {
                        let x = (position.x/10.0)as usize + PML_THICKNESS;
                        let y = (position.y/10.0)as usize + PML_THICKNESS;
                        grid.eps[x][y] = 1000.0;
                    }
                }
                UserState::SelectFreeSpace => {
                    if position != vec2(-1., -1.) {
                        let x = (position.x/10.0)as usize + PML_THICKNESS;
                        let y = (position.y/10.0)as usize + PML_THICKNESS;
                        grid.eps[x][y] = 1.0;
                    }
                }
                _ => {}
            }
        }
    }
}

fn keyboard_controls(user_state: Query<&mut UserState>,keys: Res<ButtonInput<KeyCode>>) {
    for mut state in user_state {
        if keys.just_pressed(KeyCode::KeyS) {
            *state = UserState::SelectPointSource;
        }
        if keys.just_pressed(KeyCode::KeyW) {
            *state = UserState::SelectWall;
        }
        if keys.just_pressed(KeyCode::KeyF) {
            *state = UserState::SelectFreeSpace;
        }
        if keys.just_pressed(KeyCode::Escape) {
            *state = UserState::SelectNothing;
        }

    }
}