#include "model.cuh"

float Model::train_batch(std::vector<Matrix>::iterator X_i, std::vector<Matrix>::iterator Y_i, const unsigned int batch_size, const float lr, const float momentum)
{
	float loss = 0.0f;
	for (unsigned int i = 0; i < batch_size; i++)
	{
		// create the inputs
		Matrix tmp = *X_i;

		// feedforward through the neural network
		for (auto layer = layers.begin(); layer != layers.end(); ++layer)
		{
			(*layer)->feedforward_update(tmp);
			tmp = (*layer)->neurons;
		}

		// calculate the loss
		loss += loss_function.func(*Y_i, tmp);

		// calculate the derivative for the loss function
		// reuse temporary variable, as it is no longer needed
		tmp = loss_function.deriv(*Y_i, tmp);

		// now backpropogate through the layers
		for (auto layer = layers.rbegin(); layer != layers.rend(); ++layer)
			tmp = (*layer)->backpropogate(tmp);

		++X_i; ++Y_i; // advance the iterators
	}

	for (auto layer = layers.rbegin(); layer != layers.rend(); ++layer)
		(*layer)->gradient_update(lr, momentum);

	return loss / (float) batch_size;
}

Model::Model(Loss_Function loss_function)
{
	this->loss_function = loss_function;
}

void Model::operator+=(Layer* layer)
{
	//if (layers.empty() or layers.back()->get_out_shape() == layer.get_in_shape())
		layers.push_back(layer);
	//else
		//yeet "This layer's input shape is not compatible";
}

Matrix Model::feed (const Matrix& input) const
{
	Matrix tmp = (*(layers.begin()))->feedforward(input);

	// feedforward through the neural network
	auto layer = layers.begin();
	for (++layer; layer != layers.end(); ++layer)
		tmp = (*layer)->feedforward(tmp);
	return tmp;
}
	.file	"tmpxft_00232e28_00000000-6_model.cudafe1.cpp"
	.text
#APP
	.globl _ZSt21ios_base_library_initv

#NO_APP
	.type	_ZL26__cudaUnregisterBinaryUtilv, @function
_ZL26__cudaUnregisterBinaryUtilv:
.LFB4239:
	.cfi_startproc
	endbr64
	movq	_ZL20__cudaFatCubinHandle(%rip), %rdi
	jmp	__cudaUnregisterFatBinary@PLT
	.cfi_endproc
.LFE4239:
	.size	_ZL26__cudaUnregisterBinaryUtilv, .-_ZL26__cudaUnregisterBinaryUtilv
	.align 2
	.globl	_ZN5Model11train_batchEN9__gnu_cxx17__normal_iteratorIP6MatrixSt6vectorIS2_SaIS2_EEEES7_jff
	.type	_ZN5Model11train_batchEN9__gnu_cxx17__normal_iteratorIP6MatrixSt6vectorIS2_SaIS2_EEEES7_jff, @function
_ZN5Model11train_batchEN9__gnu_cxx17__normal_iteratorIP6MatrixSt6vectorIS2_SaIS2_EEEES7_jff:
.LFB4216:
	.cfi_startproc
	.cfi_personality 0x9b,DW.ref.__gxx_personality_v0
	.cfi_lsda 0x1b,.LLSDA4216
	endbr64
	pushq	%r15
	.cfi_def_cfa_offset 16
	.cfi_offset 15, -16
	movq	%rsi, %r15
	pushq	%r14
	.cfi_def_cfa_offset 24
	.cfi_offset 14, -24
	pushq	%r13
	.cfi_def_cfa_offset 32
	.cfi_offset 13, -32
	movl	%ecx, %r13d
	pushq	%r12
	.cfi_def_cfa_offset 40
	.cfi_offset 12, -40
	pushq	%rbp
	.cfi_def_cfa_offset 48
	.cfi_offset 6, -48
	xorl	%ebp, %ebp
	pushq	%rbx
	.cfi_def_cfa_offset 56
	.cfi_offset 3, -56
	movq	%rdi, %rbx
	subq	$136, %rsp
	.cfi_def_cfa_offset 192
	movq	%rdx, 40(%rsp)
	movq	%fs:40, %rax
	movq	%rax, 120(%rsp)
	xorl	%eax, %eax
	movq	%r13, %rax
	movl	$0x00000000, 20(%rsp)
	salq	$5, %rax
	movq	%rax, 32(%rsp)
	leaq	56(%rsp), %rax
	movq	%rax, 8(%rsp)
.L3:
	movq	40(%rsp), %rax
	leaq	(%r15,%rbp), %rsi
	addq	%rbp, %rax
	movq	%rax, 24(%rsp)
	movq	32(%rsp), %rax
	cmpq	%rax, %rbp
	je	.L18
	movq	8(%rsp), %rdi
.LEHB0:
	call	_ZN6MatrixC1ERKS_@PLT
.LEHE0:
	movq	(%rbx), %r12
.L4:
	movq	%rbx, %r14
	cmpq	%rbx, %r12
	je	.L19
	movq	16(%r12), %rdi
	movq	8(%rsp), %rsi
	movq	(%rdi), %rax
.LEHB1:
	call	*8(%rax)
	movq	16(%r12), %rax
	movq	(%r12), %r12
	movups	16(%rax), %xmm2
	movups	32(%rax), %xmm3
	movups	%xmm2, 56(%rsp)
	movups	%xmm3, 72(%rsp)
	jmp	.L4
.L19:
	movq	8(%rsp), %rsi
	movq	24(%rsp), %rdi
	call	*24(%rbx)
	addss	20(%rsp), %xmm0
	leaq	88(%rsp), %r12
	movq	8(%rsp), %rdx
	movq	24(%rsp), %rsi
	movq	%r12, %rdi
	movss	%xmm0, 20(%rsp)
	call	*32(%rbx)
	movups	88(%rsp), %xmm5
	movups	104(%rsp), %xmm6
	movq	%r12, %rdi
	movups	%xmm5, 56(%rsp)
	movups	%xmm6, 72(%rsp)
	call	_ZN6MatrixD1Ev@PLT
.L6:
	cmpq	(%rbx), %r14
	je	.L20
	movq	8(%r14), %rax
	movq	8(%rsp), %rdx
	movq	%r12, %rdi
	movq	16(%rax), %rsi
	movq	(%rsi), %rax
	call	*16(%rax)
.LEHE1:
	movups	88(%rsp), %xmm7
	movups	104(%rsp), %xmm1
	movq	%r12, %rdi
	movups	%xmm7, 56(%rsp)
	movups	%xmm1, 72(%rsp)
	call	_ZN6MatrixD1Ev@PLT
	movq	8(%r14), %r14
	jmp	.L6
.L20:
	movq	8(%rsp), %rdi
	addq	$32, %rbp
	call	_ZN6MatrixD1Ev@PLT
	jmp	.L3
.L18:
	cvtsi2ssq	%r13, %xmm1
	movss	20(%rsp), %xmm0
	divss	%xmm1, %xmm0
	movq	120(%rsp), %rax
	subq	%fs:40, %rax
	je	.L13
	jmp	.L15
.L14:
	endbr64
	movq	%rax, %rbx
.L11:
	movq	8(%rsp), %rdi
	call	_ZN6MatrixD1Ev@PLT
	movq	120(%rsp), %rax
	subq	%fs:40, %rax
	jne	.L15
	movq	%rbx, %rdi
.LEHB2:
	call	_Unwind_Resume@PLT
.LEHE2:
.L15:
	call	__stack_chk_fail@PLT
.L13:
	addq	$136, %rsp
	.cfi_def_cfa_offset 56
	popq	%rbx
	.cfi_def_cfa_offset 48
	popq	%rbp
	.cfi_def_cfa_offset 40
	popq	%r12
	.cfi_def_cfa_offset 32
	popq	%r13
	.cfi_def_cfa_offset 24
	popq	%r14
	.cfi_def_cfa_offset 16
	popq	%r15
	.cfi_def_cfa_offset 8
	ret
	.cfi_endproc
.LFE4216:
	.globl	__gxx_personality_v0
	.section	.gcc_except_table,"a",@progbits
.LLSDA4216:
	.byte	0xff
	.byte	0xff
	.byte	0x1
	.uleb128 .LLSDACSE4216-.LLSDACSB4216
.LLSDACSB4216:
	.uleb128 .LEHB0-.LFB4216
	.uleb128 .LEHE0-.LEHB0
	.uleb128 0
	.uleb128 0
	.uleb128 .LEHB1-.LFB4216
	.uleb128 .LEHE1-.LEHB1
	.uleb128 .L14-.LFB4216
	.uleb128 0
	.uleb128 .LEHB2-.LFB4216
	.uleb128 .LEHE2-.LEHB2
	.uleb128 0
	.uleb128 0
.LLSDACSE4216:
	.text
	.size	_ZN5Model11train_batchEN9__gnu_cxx17__normal_iteratorIP6MatrixSt6vectorIS2_SaIS2_EEEES7_jff, .-_ZN5Model11train_batchEN9__gnu_cxx17__normal_iteratorIP6MatrixSt6vectorIS2_SaIS2_EEEES7_jff
	.align 2
	.globl	_ZN5ModelC2E13Loss_Function
	.type	_ZN5ModelC2E13Loss_Function, @function
_ZN5ModelC2E13Loss_Function:
.LFB4233:
	.cfi_startproc
	endbr64
	xorl	%eax, %eax
	movq	%rdi, 8(%rdi)
	movq	%rdi, (%rdi)
	movq	%rax, 16(%rdi)
	movq	%rsi, 24(%rdi)
	movq	%rdx, 32(%rdi)
	ret
	.cfi_endproc
.LFE4233:
	.size	_ZN5ModelC2E13Loss_Function, .-_ZN5ModelC2E13Loss_Function
	.globl	_ZN5ModelC1E13Loss_Function
	.set	_ZN5ModelC1E13Loss_Function,_ZN5ModelC2E13Loss_Function
	.align 2
	.globl	_ZN5ModelpLEP5Layer
	.type	_ZN5ModelpLEP5Layer, @function
_ZN5ModelpLEP5Layer:
.LFB4235:
	.cfi_startproc
	endbr64
	pushq	%rbp
	.cfi_def_cfa_offset 16
	.cfi_offset 6, -16
	movq	%rsi, %rbp
	pushq	%rbx
	.cfi_def_cfa_offset 24
	.cfi_offset 3, -24
	movq	%rdi, %rbx
	movl	$24, %edi
	pushq	%rax
	.cfi_def_cfa_offset 32
	call	_Znwm@PLT
	movq	%rbx, %rsi
	movq	%rbp, 16(%rax)
	movq	%rax, %rdi
	call	_ZNSt8__detail15_List_node_base7_M_hookEPS0_@PLT
	incq	16(%rbx)
	popq	%rdx
	.cfi_def_cfa_offset 24
	popq	%rbx
	.cfi_def_cfa_offset 16
	popq	%rbp
	.cfi_def_cfa_offset 8
	ret
	.cfi_endproc
.LFE4235:
	.size	_ZN5ModelpLEP5Layer, .-_ZN5ModelpLEP5Layer
	.align 2
	.globl	_ZNK5Model4feedERK6Matrix
	.type	_ZNK5Model4feedERK6Matrix, @function
_ZNK5Model4feedERK6Matrix:
.LFB4236:
	.cfi_startproc
	.cfi_personality 0x9b,DW.ref.__gxx_personality_v0
	.cfi_lsda 0x1b,.LLSDA4236
	endbr64
	pushq	%r13
	.cfi_def_cfa_offset 16
	.cfi_offset 13, -16
	pushq	%r12
	.cfi_def_cfa_offset 24
	.cfi_offset 12, -24
	pushq	%rbp
	.cfi_def_cfa_offset 32
	.cfi_offset 6, -32
	movq	%rsi, %rbp
	pushq	%rbx
	.cfi_def_cfa_offset 40
	.cfi_offset 3, -40
	movq	%rdi, %rbx
	subq	$56, %rsp
	.cfi_def_cfa_offset 96
	movq	%fs:40, %rax
	movq	%rax, 40(%rsp)
	movq	(%rsi), %rax
	leaq	8(%rsp), %r13
	movq	16(%rax), %rsi
	movq	(%rsi), %rax
.LEHB3:
	call	*(%rax)
.LEHE3:
	movq	0(%rbp), %rax
	movq	(%rax), %r12
.L25:
	cmpq	%rbp, %r12
	je	.L34
	movq	16(%r12), %rsi
	movq	%rbx, %rdx
	movq	%r13, %rdi
	movq	(%rsi), %rax
.LEHB4:
	call	*(%rax)
.LEHE4:
	movups	8(%rsp), %xmm0
	movups	24(%rsp), %xmm1
	movq	%r13, %rdi
	movups	%xmm0, (%rbx)
	movups	%xmm1, 16(%rbx)
	call	_ZN6MatrixD1Ev@PLT
	movq	(%r12), %r12
	jmp	.L25
.L30:
	endbr64
	movq	%rax, %rbp
.L27:
	movq	%rbx, %rdi
	call	_ZN6MatrixD1Ev@PLT
	movq	40(%rsp), %rax
	subq	%fs:40, %rax
	jne	.L33
	movq	%rbp, %rdi
.LEHB5:
	call	_Unwind_Resume@PLT
.LEHE5:
.L34:
	movq	40(%rsp), %rax
	subq	%fs:40, %rax
	je	.L29
.L33:
	call	__stack_chk_fail@PLT
.L29:
	addq	$56, %rsp
	.cfi_def_cfa_offset 40
	movq	%rbx, %rax
	popq	%rbx
	.cfi_def_cfa_offset 32
	popq	%rbp
	.cfi_def_cfa_offset 24
	popq	%r12
	.cfi_def_cfa_offset 16
	popq	%r13
	.cfi_def_cfa_offset 8
	ret
	.cfi_endproc
.LFE4236:
	.section	.gcc_except_table
.LLSDA4236:
	.byte	0xff
	.byte	0xff
	.byte	0x1
	.uleb128 .LLSDACSE4236-.LLSDACSB4236
.LLSDACSB4236:
	.uleb128 .LEHB3-.LFB4236
	.uleb128 .LEHE3-.LEHB3
	.uleb128 0
	.uleb128 0
	.uleb128 .LEHB4-.LFB4236
	.uleb128 .LEHE4-.LEHB4
	.uleb128 .L30-.LFB4236
	.uleb128 0
	.uleb128 .LEHB5-.LFB4236
	.uleb128 .LEHE5-.LEHB5
	.uleb128 0
	.uleb128 0
.LLSDACSE4236:
	.text
	.size	_ZNK5Model4feedERK6Matrix, .-_ZNK5Model4feedERK6Matrix
	.section	.text.startup,"ax",@progbits
	.type	_ZL24__sti____cudaRegisterAllv, @function
_ZL24__sti____cudaRegisterAllv:
.LFB4262:
	.cfi_startproc
	endbr64
	pushq	%rax
	.cfi_def_cfa_offset 16
	leaq	_ZL15__fatDeviceText(%rip), %rdi
	call	__cudaRegisterFatBinary@PLT
	movq	%rax, %rdi
	movq	%rax, _ZL20__cudaFatCubinHandle(%rip)
	call	__cudaRegisterFatBinaryEnd@PLT
	leaq	_ZL26__cudaUnregisterBinaryUtilv(%rip), %rdi
	popq	%rdx
	.cfi_def_cfa_offset 8
	jmp	atexit@PLT
	.cfi_endproc
.LFE4262:
	.size	_ZL24__sti____cudaRegisterAllv, .-_ZL24__sti____cudaRegisterAllv
	.section	.init_array,"aw"
	.align 8
	.quad	_ZL24__sti____cudaRegisterAllv
	.section	.nvFatBinSegment,"aw"
	.align 8
	.type	_ZL15__fatDeviceText, @object
	.size	_ZL15__fatDeviceText, 24
_ZL15__fatDeviceText:
	.long	1180844977
	.long	1
	.quad	fatbinData
	.quad	0
	.local	_ZL20__cudaFatCubinHandle
	.comm	_ZL20__cudaFatCubinHandle,8,8
	.hidden	DW.ref.__gxx_personality_v0
	.weak	DW.ref.__gxx_personality_v0
	.section	.data.rel.local.DW.ref.__gxx_personality_v0,"awG",@progbits,DW.ref.__gxx_personality_v0,comdat
	.align 8
	.type	DW.ref.__gxx_personality_v0, @object
	.size	DW.ref.__gxx_personality_v0, 8
DW.ref.__gxx_personality_v0:
	.quad	__gxx_personality_v0
	.ident	"GCC: (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0"
	.section	.note.GNU-stack,"",@progbits
	.section	.note.gnu.property,"a"
	.align 8
	.long	1f - 0f
	.long	4f - 1f
	.long	5
0:
	.string	"GNU"
1:
	.align 8
	.long	0xc0000002
	.long	3f - 2f
2:
	.long	0x3
3:
	.align 8
4:

#pragma once
#include <hip/hip_runtime.h>


#include "layer.cuh"
#include "dense.cuh"
#include "losses.cuh"

#include <iterator>
#include <list>
#include <vector>

class Model
{
	std::list<Layer*> layers;
	Loss_Function loss_function;


public:
	float train_batch(std::vector<Matrix>::iterator, std::vector<Matrix>::iterator, const unsigned int, const float, const float=0.0f); // train one batch and return the average loss on that batch

	Model(Loss_Function);

	void operator+= (Layer*); // add a new layer to the neural network
	Matrix feed (const Matrix&) const; // use the model to predict from a given input

	void train(const std::vector<Matrix>&, const std::vector<Matrix>&, const unsigned int batch_size);
};
	.text
	.file	"model.hip"
                                        # Start of file scope inline assembly
	.globl	_ZSt21ios_base_library_initv

                                        # End of file scope inline assembly
	.type	__hip_cuid_,@object             # @__hip_cuid_
	.bss
	.globl	__hip_cuid_
__hip_cuid_:
	.byte	0                               # 0x0
	.size	__hip_cuid_, 1

	.ident	"AMD clang version 18.0.0git (https://github.com/RadeonOpenCompute/llvm-project roc-6.3.2 25012 e5bf7e55c91490b07c49d8960fa7983d864936c4)"
	.section	".note.GNU-stack","",@progbits
	.addrsig
	.addrsig_sym __hip_cuid_

	.text
	.p2alignl 7, 3214868480
	.fill 96, 4, 3214868480
	.type	__hip_cuid_,@object             ; @__hip_cuid_
	.section	.bss,"aw",@nobits
	.globl	__hip_cuid_
__hip_cuid_:
	.byte	0                               ; 0x0
	.size	__hip_cuid_, 1

	.ident	"AMD clang version 18.0.0git (https://github.com/RadeonOpenCompute/llvm-project roc-6.3.2 25012 e5bf7e55c91490b07c49d8960fa7983d864936c4)"
	.section	".note.GNU-stack","",@progbits
	.addrsig
	.addrsig_sym __hip_cuid_
	.amdgcn_target "amdgcn-amd-amdhsa--gfx1100"
	.amdgpu_metadata
---
amdhsa.kernels:  []
amdhsa.target:   amdgcn-amd-amdhsa--gfx1100
amdhsa.version:
  - 1
  - 2
...

	.end_amdgpu_metadata
