-module(mlx_dist_worker).
-behaviour(gen_server).

%% API
-export([start_link/0, start_link/1,
         register_with_coordinator/1,
         initialize_training/2,
         set_data/2,
         compute_gradients/0,
         get_gradients/0,
         update_parameters/1,
         stop_training/0,
         get_status/0]).

%% gen_server callbacks
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
         terminate/2, code_change/3]).

-record(state, {
    coordinator = undefined,
    model = undefined,         % Current model (MLX arrays)
    data = undefined,          % Training data
    batch_size = 32,
    current_batch = 0,
    gradients = undefined,     % Latest computed gradients
    optimizer_state = undefined,
    training_config = undefined,
    status = idle,
    device = gpu,              % gpu or cpu
    stats = #{}
}).

%%====================================================================
%% API
%%====================================================================

start_link() ->
    start_link([]).

start_link(Options) ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, Options, []).

%% Register this worker with a coordinator
register_with_coordinator(CoordinatorNode) ->
    gen_server:call(?MODULE, {register_coordinator, CoordinatorNode}).

%% Initialize model and training configuration
initialize_training(ModelConfig, TrainingConfig) ->
    gen_server:call(?MODULE, {initialize_training, ModelConfig, TrainingConfig}).

%% Set training data for this worker
set_data(Data, BatchIndex) ->
    gen_server:call(?MODULE, {set_data, Data, BatchIndex}).

%% Compute gradients for current batch
compute_gradients() ->
    gen_server:cast(?MODULE, compute_gradients).

%% Get computed gradients
get_gradients() ->
    gen_server:call(?MODULE, get_gradients, 30000).

%% Update model parameters
update_parameters(Parameters) ->
    gen_server:call(?MODULE, {update_parameters, Parameters}).

%% Stop training
stop_training() ->
    gen_server:call(?MODULE, stop_training).

%% Get worker status
get_status() ->
    gen_server:call(?MODULE, get_status).

%%====================================================================
%% gen_server callbacks
%%====================================================================

init(Options) ->
    io:format("MLX Distributed Worker started on ~p~n", [node()]),
    
    %% Start MLX application
    application:start(mlx),
    
    %% Set device
    Device = proplists:get_value(device, Options, gpu),
    case Device of
        gpu -> mlx:use_gpu();
        cpu -> mlx:use_cpu()
    end,
    
    {ok, #state{device = Device}}.

handle_call({register_coordinator, CoordinatorNode}, _From, State) ->
    io:format("Registering with coordinator: ~p~n", [CoordinatorNode]),
    {reply, ok, State#state{coordinator = CoordinatorNode}};

handle_call({initialize_training, ModelConfig, TrainingConfig}, _From, State) ->
    try
        %% Initialize model based on config
        Model = initialize_model(ModelConfig),
        
        %% Initialize optimizer state
        OptimizerState = initialize_optimizer(TrainingConfig),
        
        BatchSize = maps:get(batch_size, TrainingConfig, 32),
        
        NewState = State#state{
            model = Model,
            training_config = TrainingConfig,
            optimizer_state = OptimizerState,
            batch_size = BatchSize,
            status = initialized
        },
        
        io:format("Worker initialized with model on ~p~n", [State#state.device]),
        {reply, ok, NewState}
    catch
        Error:Reason ->
            io:format("Initialization error: ~p:~p~n", [Error, Reason]),
            {reply, {error, Reason}, State}
    end;

handle_call({set_data, Data, BatchIndex}, _From, State) ->
    io:format("Worker received ~p data samples~n", [length(Data)]),
    {reply, ok, State#state{data = Data, current_batch = BatchIndex}};

handle_call(get_gradients, _From, State) ->
    case State#state.gradients of
        undefined ->
            {reply, {error, no_gradients}, State};
        Gradients ->
            {reply, {ok, Gradients}, State}
    end;

handle_call({update_parameters, Parameters}, _From, State) ->
    %% Update model with new parameters
    NewModel = update_model_parameters(State#state.model, Parameters),
    {reply, ok, State#state{model = NewModel}};

handle_call(stop_training, _From, State) ->
    {reply, ok, State#state{status = idle}};

handle_call(get_status, _From, State) ->
    Status = #{
        node => node(),
        device => State#state.device,
        status => State#state.status,
        has_data => State#state.data =/= undefined,
        has_model => State#state.model =/= undefined,
        stats => State#state.stats
    },
    {reply, Status, State};

handle_call(_Request, _From, State) ->
    {reply, {error, unknown_request}, State}.

handle_cast(compute_gradients, State) ->
    %% Compute gradients asynchronously
    case State#state.status of
        initialized ->
            NewState = compute_gradients_internal(State),
            {noreply, NewState};
        _ ->
            {noreply, State}
    end;

handle_cast(_Msg, State) ->
    {noreply, State}.

handle_info(_Info, State) ->
    {noreply, State}.

terminate(_Reason, _State) ->
    ok.

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%%====================================================================
%% Internal functions
%%====================================================================

initialize_model(ModelConfig) ->
    %% Example: Initialize a simple neural network
    case maps:get(type, ModelConfig, simple_nn) of
        simple_nn ->
            InputSize = maps:get(input_size, ModelConfig, 784),
            HiddenSize = maps:get(hidden_size, ModelConfig, 256),
            OutputSize = maps:get(output_size, ModelConfig, 10),
            
            %% Initialize weights with Xavier initialization
            W1 = xavier_init([InputSize, HiddenSize]),
            B1 = mlx:zeros([HiddenSize]),
            W2 = xavier_init([HiddenSize, OutputSize]),
            B2 = mlx:zeros([OutputSize]),
            
            #{w1 => W1, b1 => B1, w2 => W2, b2 => B2};
        
        cnn ->
            %% TODO: Initialize CNN architecture
            #{};
        
        transformer ->
            %% TODO: Initialize transformer architecture
            #{}
    end.

xavier_init(Shape) ->
    %% Xavier/Glorot initialization
    [InputDim, OutputDim] = Shape,
    Scale = math:sqrt(2.0 / (InputDim + OutputDim)),
    
    %% Create random matrix and scale
    Random = mlx:subtract(
        mlx:multiply(mlx:array(create_random_matrix(InputDim, OutputDim)), 
                     mlx:array(2)),
        mlx:array(1)
    ),
    mlx:multiply(Random, mlx:array(Scale)).

create_random_matrix(Rows, Cols) ->
    [[rand:uniform() || _ <- lists:seq(1, Cols)] || _ <- lists:seq(1, Rows)].

initialize_optimizer(TrainingConfig) ->
    %% Initialize optimizer state (e.g., Adam momentum terms)
    LearningRate = maps:get(learning_rate, TrainingConfig, 0.001),
    #{
        learning_rate => LearningRate,
        iteration => 0,
        momentum => #{},
        velocity => #{}
    }.

compute_gradients_internal(State) ->
    case {State#state.model, State#state.data} of
        {undefined, _} ->
            State;
        {_, undefined} ->
            State;
        {Model, Data} ->
            %% Get next batch
            {BatchX, BatchY} = get_next_batch(Data, State#state.current_batch, 
                                             State#state.batch_size),
            
            %% Forward pass
            {Output, Intermediates} = forward_pass(Model, BatchX),
            
            %% Compute loss
            Loss = cross_entropy_loss(Output, BatchY),
            
            %% Backward pass (compute gradients)
            Gradients = backward_pass(Model, Intermediates, Output, BatchY),
            
            %% Update stats
            Stats = update_stats(State#state.stats, Loss),
            
            %% Update batch index
            NextBatch = (State#state.current_batch + 1) rem 
                       (length(Data) div State#state.batch_size),
            
            State#state{
                gradients = Gradients,
                current_batch = NextBatch,
                stats = Stats
            }
    end.

get_next_batch(Data, BatchIndex, BatchSize) ->
    %% Extract batch from data
    %% Assuming Data is a list of {input, label} tuples
    StartIdx = BatchIndex * BatchSize + 1,
    EndIdx = erlang:min(StartIdx + BatchSize - 1, length(Data)),
    
    Batch = lists:sublist(Data, StartIdx, EndIdx - StartIdx + 1),
    
    %% Separate inputs and labels
    {Inputs, Labels} = lists:unzip(Batch),
    
    %% Convert to MLX arrays
    BatchX = mlx:array(Inputs),
    BatchY = mlx:array(Labels),
    
    {BatchX, BatchY}.

forward_pass(#{w1 := W1, b1 := B1, w2 := W2, b2 := B2}, Input) ->
    %% Simple 2-layer neural network forward pass
    
    %% Layer 1: Linear + ReLU
    Z1 = mlx:add(mlx:matmul(Input, W1), B1),
    A1 = relu(Z1),
    
    %% Layer 2: Linear
    Z2 = mlx:add(mlx:matmul(A1, W2), B2),
    
    %% Output with softmax
    Output = mlx:softmax(Z2, 1),
    
    %% Return output and intermediates for backward pass
    {Output, #{input => Input, z1 => Z1, a1 => A1, z2 => Z2}}.

relu(X) ->
    %% ReLU activation: max(0, x)
    mlx:maximum(X, mlx:array(0)).

cross_entropy_loss(Predictions, Labels) ->
    %% Cross-entropy loss for classification
    %% -sum(y * log(pred)) / batch_size
    
    %% One-hot encode labels if needed
    OneHot = one_hot_encode(Labels, mlx:shape(Predictions)),
    
    %% Compute loss
    LogPreds = mlx:log(mlx:add(Predictions, mlx:array(0.0000001))), % Add epsilon for stability
    Loss = mlx:negative(mlx:sum(mlx:multiply(OneHot, LogPreds))),
    
    %% Average over batch
    BatchSize = element(1, mlx:shape(Predictions)),
    mlx:divide(Loss, mlx:array(BatchSize)).

one_hot_encode(Labels, {BatchSize, NumClasses}) ->
    %% Convert labels to one-hot encoding
    %% This is simplified - in practice would need proper implementation
    mlx:eye(NumClasses).

backward_pass(Model = #{w1 := W1, b1 := B1, w2 := W2, b2 := B2}, 
              Intermediates, Output, Labels) ->
    %% Backward pass to compute gradients
    
    #{input := Input, z1 := Z1, a1 := A1, z2 := Z2} = Intermediates,
    
    %% Output layer gradients
    DZ2 = mlx:subtract(Output, one_hot_encode(Labels, mlx:shape(Output))),
    DW2 = mlx:matmul(mlx:transpose(A1), DZ2),
    DB2 = mlx:sum(DZ2, 0),
    
    %% Hidden layer gradients
    DA1 = mlx:matmul(DZ2, mlx:transpose(W2)),
    DZ1 = mlx:multiply(DA1, relu_derivative(Z1)),
    DW1 = mlx:matmul(mlx:transpose(Input), DZ1),
    DB1 = mlx:sum(DZ1, 0),
    
    %% Return gradients
    #{dw1 => DW1, db1 => DB1, dw2 => DW2, db2 => DB2}.

relu_derivative(X) ->
    %% Derivative of ReLU: 1 if x > 0, else 0
    mlx:greater(X, mlx:array(0)).

update_model_parameters(Model, Parameters) ->
    %% Update model with new parameters
    %% Parameters should match model structure
    case {Model, Parameters} of
        {undefined, _} -> undefined;
        {_, undefined} -> Model;
        _ ->
            %% Assuming parameters is a map with same structure as model
            maps:merge(Model, Parameters)
    end.

update_stats(Stats, Loss) ->
    %% Update training statistics
    LossValue = mlx:to_list(Loss),
    
    Stats#{
        last_loss => LossValue,
        loss_history => [LossValue | maps:get(loss_history, Stats, [])]
    }.

min(A, B) when A < B -> A;
min(_, B) -> B.