#include <string.h>
#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <time.h>

#define TONE_CODE     1
#define BARRAGE_CODE   2
#define SWEEP_CODE    3
#define PULSE_CODE    4
#define PROTOCOL_CODE 5

void print_help() {
	printf("Usage:\n");
	printf("  ./main [options]\n\n");
	printf("Options:\n");
	printf("  --method     , -m  <int>      Method type:\n");
	printf("                                1 = Tone\n");
	printf("                                2 = Barrage\n");
	printf("                                3 = Sweep\n");
	printf("                                4 = Pulse\n");
	printf("                                5 = Protocol\n");
	printf("  --path       , -o  <string>   Output file path (e.g. \"jam.iq\")\n");
	printf("  --sample_rate, -s  <double>   Sample rate in Hz (max: 2e7)\n");
	printf("  --total_time , -t  <double>   Total duration in seconds\n");
	printf("  --amplitude  , -a  <double>   Signal amplitude (range: < -127, +127 >)\n");
	printf("  --frequency  , -f  <double>   Tone frequency in Hz (range: < -sample_rate/2, +sample_rate/2 >)\n");
	printf("  --tone_count , -n  <int>      Number of tones\n");
	printf("  --bandwidth  , -b  <double>   Signal bandwidth in Hz (max: sample_rate)\n");
	printf("  --sweep_time , -w  <double>   Time for a full sweep, in seconds\n");
	printf("                                (this will override total_time)\n");
	printf("  --duty_cycle , -d  <double>   Duty cycle as a fraction (range: < 0, 1>)\n");
	printf("  --period     , -p  <double>   Period in seconds\n");
	printf("  --bit_rate   , -r  <double>   Bit rate in bits per second (max: 1e7)\n");
	printf("\nExamples:\n");
	printf("  ./main --frequency 1e6 -a 127 -o \"jam.iq\" -m 2\n");
	printf("  ./main -m 3 -s 2e6 -t 10 --path \"sweep.iq\"\n");
	printf("\nNotes:\n");
	printf("  - All arguments with numerical values expect floating point numbers unless stated.\n");
	printf("  - 'method' must be an integer from 1 to 5 indicating the signal generation strategy.\n");
	printf("  - Both long (--argument) and short (-x) flags are accepted.\n");
	printf("  - Arguments can appear in any order.\n");
}

typedef struct {
	int method;
	char* path;
	double sample_rate;
	double total_time;
	double amplitude;
	double frequency;
	double tone_count;
	double bandwidth;
	double sweep_time;
	double duty_cycle;
	double period;
	double bit_rate;
} Args;

Args parse_args(int argc, char* argv[]) {
	if (argc == 1) {
		print_help();
		exit(0);
	}

	char *default_path = "jam.iq";
	Args args = {
		.method      = 1,
		.path        = default_path,
		.sample_rate = 02e+7,
		.total_time  = 0.01,
		.amplitude   = 127,
		.frequency   = 0,
		.tone_count  = 5,
		.bandwidth   = 01e+6,
		.sweep_time  = 0.01,
		.duty_cycle  = 24e-6,
		.period      = 05e-2,
		.bit_rate    = 01e+5
	};

	for (int i = 1; i < argc; i++) {
		char* arg = argv[i];

		if (strcmp(arg, "--help") == 0 || strcmp(arg, "-h") == 0) {
			print_help();
			exit(0);
		} else if ((strcmp(arg, "--method") == 0 || strcmp(arg, "-m") == 0) && i + 1 < argc) {
			args.method = atof(argv[++i]);
		} else if ((strcmp(arg, "--path") == 0 || strcmp(arg, "-o") == 0) && i + 1 < argc) {
			args.path = argv[++i];
		} else if ((strcmp(arg, "--sample_rate") == 0 || strcmp(arg, "-s") == 0) && i + 1 < argc) {
			args.sample_rate = atof(argv[++i]);
		} else if ((strcmp(arg, "--total_time") == 0 || strcmp(arg, "-t") == 0) && i + 1 < argc) {
			args.total_time = atof(argv[++i]);
		} else if ((strcmp(arg, "--amplitude") == 0 || strcmp(arg, "-a") == 0) && i + 1 < argc) {
			args.amplitude = atof(argv[++i]);
		} else if ((strcmp(arg, "--tone_count") == 0 || strcmp(arg, "-n") == 0) && i + 1 < argc) {
			args.tone_count = atof(argv[++i]);
		} else if ((strcmp(arg, "--bandwidth") == 0 || strcmp(arg, "-b") == 0) && i + 1 < argc) {
			args.bandwidth = atof(argv[++i]);
		} else if ((strcmp(arg, "--bandwidth") == 0 || strcmp(arg, "-b") == 0) && i + 1 < argc) {
			args.bandwidth = atof(argv[++i]);
		} else if ((strcmp(arg, "--sweep_time") == 0 || strcmp(arg, "-w") == 0) && i + 1 < argc) {
			args.sweep_time = atof(argv[++i]);
		} else if ((strcmp(arg, "--duty_cycle") == 0 || strcmp(arg, "-d") == 0) && i + 1 < argc) {
			args.duty_cycle = atof(argv[++i]);
		} else if ((strcmp(arg, "--period") == 0 || strcmp(arg, "-p") == 0) && i + 1 < argc) {
			args.period = atof(argv[++i]);
		} else if ((strcmp(arg, "--bit_rate") == 0 || strcmp(arg, "-r") == 0) && i + 1 < argc) {
			args.bit_rate = atof(argv[++i]);
		} else {
			fprintf(stderr, "Error: Unrecognized or incomplete flag '%s'\n", arg);
			exit(0);
		}
	}
	if (args.method > 6 || args.method < 1) {
		fprintf(stderr, "Error: Method must be in range from 1 to 5!\n");
		exit(0);
	} else if (args.sample_rate < 0) {
		fprintf(stderr, "Error: Sample rate must be positive!\n");
		exit(0);
	} else if (args.total_time < 0) {
		fprintf(stderr, "Error: Total time must be positive!\n");
		exit(0);
	} else if (args.amplitude < -127 || args.amplitude > 127) {
		fprintf(stderr, "Error: Amplitude must be in range from -127 to +127!\n");
		exit(0);
	} else if (args.frequency < -args.sample_rate/2 || args.frequency > args.sample_rate/2) {
		fprintf(stderr, "Error: frequency must be in range of +- half the sample_rate!\n");
		exit(0);
	} else if (args.tone_count <= 1) {
		fprintf(stderr, "Error: Tone count must be greater than one!\n");
		exit(0);
	} else if (args.bandwidth < 0) {
		fprintf(stderr, "Error: Bandwidth must be positive!\n");
		exit(0);
	} else if (args.sweep_time < 0) {
		fprintf(stderr, "Error: Sweep time must be positive!\n");
		exit(0);
	} else if (args.duty_cycle > 1 || args.duty_cycle < 0) {
		fprintf(stderr, "Error: Duty cycle must be in range from 0 to 1!\n");
		exit(0);
	} else if (args.period < 0) {
		fprintf(stderr, "Error: Period must be positive!\n");
		exit(0);
	} else if (args.bit_rate < 0) {
		fprintf(stderr, "Error: Bit rate must be positive!\n");
		exit(0);
	}
	return args;
}

double rect(double time, double tau) {
	return (fabs(time) < (tau / 2)) ? 1.0 : 0.0;
}

int main(int argc, char *argv[]) {
	Args args = parse_args(argc, argv);

	FILE *file = fopen(args.path, "wb");
	if (!file) {
		printf("Error: Failed to open file\n");
		exit(0);
	}

	uint32_t num_samples = args.total_time * args.sample_rate;

	int8_t *IQ_data = (int8_t *)malloc(num_samples * 2 * sizeof(int8_t));

	if (!IQ_data) {
		printf("Error: Failed to allocate IQ_data\n");
		exit(0);
	}

	uint32_t data_len = (uint32_t)((args.total_time * args.bit_rate) + 1);
	uint8_t noise_data[num_samples];
	srand(time(NULL));
	for (uint32_t i = 0; i < data_len; i++) {
		noise_data[i] = ( ( rand() % 2 ) * 2) - 1;
	}

	double time = 0;
	double freq = 0;
	double sum = 0;

	switch (args.method) {
		case TONE_CODE:
			for (uint32_t sample = 0; sample < num_samples; ++sample) {
				time = (double)sample/args.sample_rate;
				IQ_data[2 * sample + 0] = (int8_t)(args.amplitude * cos(2.0 * M_PI * args.frequency * time));
				IQ_data[2 * sample + 1] = (int8_t)(args.amplitude * sin(2.0 * M_PI * args.frequency * time));
			}
			break;
		case BARRAGE_CODE:
			for (uint32_t sample = 0; sample < num_samples; ++sample) {
				time = (double)sample/args.sample_rate;
				sum = 1;
				for (int i = 1; i <= args.tone_count; ++i) {
					freq = (i * args.bandwidth) / (2 * args.tone_count);
					sum += sin(2.0 * M_PI * time * freq);
				}
				IQ_data[2 * sample + 0] = (int8_t)(0);
				IQ_data[2 * sample + 1] = (int8_t)((args.amplitude * sum) / (args.tone_count + 1));
			}
			break;
		case SWEEP_CODE:
			for (uint32_t sample = 0; sample < num_samples; ++sample) {
				time = (double)sample/args.sample_rate;
				freq = (-args.bandwidth/2.0) + (args.bandwidth/2.0 * time/args.sweep_time);
				IQ_data[2 * sample + 0] = (int8_t)(0);
				IQ_data[2 * sample + 1] = (int8_t)(args.amplitude * sin(2.0 * M_PI * freq * time));
			}
			break;
		case PULSE_CODE:
			for (uint32_t sample = 0; sample < num_samples; ++sample) {
				time = (double)sample/args.sample_rate;
				IQ_data[2 * sample + 0] = (int8_t)(0);
				IQ_data[2 * sample + 1] = (int8_t)(args.amplitude * rect(time - (int)(time/args.period) * args.period, args.period * args.duty_cycle));
			}
			break;
		case PROTOCOL_CODE:
			for (uint32_t sample = 0; sample < num_samples; ++sample) {
				time = (double)sample/args.sample_rate;
				IQ_data[2 * sample + 0] = (int8_t)(0);
				IQ_data[2 * sample + 1] = (int8_t)(args.amplitude * noise_data[(uint32_t)(time * args.bit_rate)]);
			}
			break;

		default:
			printf("Error: Unknown type of jamming!\n");
			exit(0);
	}

	fwrite(IQ_data, sizeof(int8_t), num_samples * 2, file);
	fclose(file);
	free(IQ_data);
	
	printf("Jam signal saved to %s\n", args.path);

	return 0;
}