import os
import csv
from Damage_calculation import DamageCalculator

class CSVGenerator:
    def __init__(self, output_folder):
        self.damage_folder = os.path.join(output_folder, 'damage/labels')
        self.parts_folder = os.path.join(output_folder, 'parts/labels')
        self.output_csv = os.path.join(output_folder, 'damage_estimation.csv')
        self.calculator = DamageCalculator()

    def process_files(self):
        damage_files = sorted(os.listdir(self.damage_folder))
        parts_files = sorted(os.listdir(self.parts_folder))

        total_cost_min = 0
        total_cost_max = 0
        damage_data = []

        with open(self.output_csv, 'w', newline='') as csvfile:
            fieldnames = ['Car Parts', 'Damage Type', 'Damage Count', 'Severity', 'Estimated Cost Range']
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            writer.writeheader()

            for damage_file, parts_file in zip(damage_files, parts_files):
                damage_path = os.path.join(self.damage_folder, damage_file)
                parts_path = os.path.join(self.parts_folder, parts_file)

                damage_classes, damage_polygons = self.calculator.parse_coordinates(damage_path)
                part_classes, part_polygons = self.calculator.parse_coordinates(parts_path)

                damage_polygons = list(zip(damage_classes, damage_polygons))
                part_polygons = list(zip(part_classes, part_polygons))

                intersection_info = self.calculator.calculate_intersection_area(damage_polygons, part_polygons)

                for info in intersection_info:
                    writer.writerow({
                        'Car Parts': info['part_class'],
                        'Damage Type': info['damage_class'],
                        'Damage Count': info['count'],
                        'Severity': info['severity'],
                        'Estimated Cost Range': f"Rs. {info['cost_min']}-{info['cost_max']}"
                    })
                    total_cost_min += info['cost_min']
                    total_cost_max += info['cost_max']

                    damage_data.append({
                        'damage_class': info['damage_class'],
                        'part_class': info['part_class'],
                        'coverage_percentage': info['coverage_percentage']
                    })

            writer.writerow({
                'Car Parts': 'Total',
                'Damage Type': '',
                'Damage Count': '',
                'Severity': '',
                'Estimated Cost Range': f"Rs. {total_cost_min}-{total_cost_max}"
            })

        return damage_data

def main(output_folder):
    generator = CSVGenerator(output_folder)
    generator.process_files()

if __name__ == "__main__":
    output_folder = 'Output'
    main(output_folder)